From 014e44e35292eba9c8c9891064aea25351383fa8 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 5 Jul 2024 19:56:35 +0530
Subject: [PATCH] feat: support `SciMLBase.remake_initializeprob`

---
 src/systems/nonlinear/initializesystem.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index e9f70c09e9..06a3e7e58b 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/src/systems/nonlinear/initializesystem.jl
@@ -107,3 +107,13 @@ function generate_initializesystem(sys::ODESystem;
 
     return sys_nl
 end
+
+function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
+    if !(eltype(u0) <: Pair) && !isempty(u0)
+        return odefn.initializeprob, odefn.initializeprobmap
+    end
+    initprob = InitializationProblem(sys, t0, u0, p)
+    display(initprob.f.sys); @show equations(initprob.f.sys)
+    initprobmap = getu(initprob, unknowns(sys))
+    return initprob, initprobmap
+end