From 64d20636b4d173ff117da8737ad975662deb0856 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@juliahub.com>
Date: Wed, 26 Jun 2024 12:59:12 +0530
Subject: [PATCH] fix: fix initialization with defaults dependent on indepvar

---
 src/systems/diffeqs/abstractodesystem.jl | 18 +++++++++++++++---
 src/systems/parameter_buffer.jl          |  7 +++++--
 test/initial_values.jl                   |  9 +++++++++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index 18ec4e928a..b33df83989 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -675,6 +675,7 @@ Take dictionaries with initial conditions and parameters and convert them to num
 function get_u0_p(sys,
         u0map,
         parammap = nothing;
+        t0 = nothing,
         use_union = true,
         tofloat = true,
         symbolic_u0 = false)
@@ -682,6 +683,9 @@ function get_u0_p(sys,
     ps = parameters(sys)
 
     defs = defaults(sys)
+    if t0 !== nothing
+        defs[get_iv(sys)] = t0
+    end
     if parammap !== nothing
         defs = mergedefaults(defs, parammap, ps)
     end
@@ -717,14 +721,19 @@ function get_u0_p(sys,
     end
     p = varmap_to_vars(parammap, ps; defaults = defs, tofloat, use_union)
     p = p === nothing ? SciMLBase.NullParameters() : p
+    t0 !== nothing && delete!(defs, get_iv(sys))
     u0, p, defs
 end
 
 function get_u0(
-        sys, u0map, parammap = nothing; symbolic_u0 = false, toterm = default_toterm)
+        sys, u0map, parammap = nothing; symbolic_u0 = false,
+        toterm = default_toterm, t0 = nothing)
     dvs = unknowns(sys)
     ps = parameters(sys)
     defs = defaults(sys)
+    if t0 !== nothing
+        defs[get_iv(sys)] = t0
+    end
     if parammap !== nothing
         defs = mergedefaults(defs, parammap, ps)
     end
@@ -745,6 +754,7 @@ function get_u0(
     else
         u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true, toterm)
     end
+    t0 !== nothing && delete!(defs, get_iv(sys))
     return u0, defs
 end
 
@@ -819,13 +829,14 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     end
 
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        u0, defs = get_u0(sys, trueinit, parammap; symbolic_u0)
+        u0, defs = get_u0(sys, trueinit, parammap; symbolic_u0,
+            t0 = constructor <: Union{DDEFunction, SDDEFunction} ? nothing : t)
         check_eqs_u0(eqs, dvs, u0; kwargs...)
         p = if parammap === nothing ||
                parammap == SciMLBase.NullParameters() && isempty(defs)
             nothing
         else
-            MTKParameters(sys, parammap, trueinit; eval_expression, eval_module)
+            MTKParameters(sys, parammap, trueinit; t0 = t, eval_expression, eval_module)
         end
     else
         u0, p, defs = get_u0_p(sys,
@@ -833,6 +844,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
             parammap;
             tofloat,
             use_union,
+            t0 = constructor <: Union{DDEFunction, SDDEFunction} ? nothing : t,
             symbolic_u0)
         p, split_idxs = split_parameters_by_type(p)
         if p isa Tuple
diff --git a/src/systems/parameter_buffer.jl b/src/systems/parameter_buffer.jl
index ebff9adf74..b599562e26 100644
--- a/src/systems/parameter_buffer.jl
+++ b/src/systems/parameter_buffer.jl
@@ -15,8 +15,8 @@ end
 
 function MTKParameters(
         sys::AbstractSystem, p, u0 = Dict(); tofloat = false, use_union = false,
-        eval_expression = false, eval_module = @__MODULE__)
-    ic::IndexCache = if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        t0 = nothing, eval_expression = false, eval_module = @__MODULE__)
+    ic = if has_index_cache(sys) && get_index_cache(sys) !== nothing
         get_index_cache(sys)
     else
         error("Cannot create MTKParameters if system does not have index_cache")
@@ -43,6 +43,9 @@ function MTKParameters(
     defs = merge(defs, u0)
     defs = merge(Dict(eq.lhs => eq.rhs for eq in observed(sys)), defs)
     bigdefs = merge(defs, p)
+    if t0 !== nothing
+        bigdefs[get_iv(sys)] = t0
+    end
     p = Dict()
     missing_params = Set()
     pdeps = has_parameter_dependencies(sys) ? parameter_dependencies(sys) : nothing
diff --git a/test/initial_values.jl b/test/initial_values.jl
index 4c2acd0aac..ce068e73b2 100644
--- a/test/initial_values.jl
+++ b/test/initial_values.jl
@@ -108,3 +108,12 @@ prob = ODEProblem(sys, [], (0.0, 1.0), [A1 => 0.3])
         @test isempty(ModelingToolkit.defaults(sys))
     end
 end
+
+# Using indepvar in initialization
+# Issue#2799
+@variables x(t)
+@parameters p
+@mtkbuild sys = ODESystem([D(x) ~ p], t; defaults = [x => t, p => 2t])
+prob = ODEProblem(structural_simplify(sys), [], (1.0, 2.0), [])
+@test prob[x] == 1.0
+@test prob.ps[p] == 2.0