From f94b26f18814ee6649da75985103fe9cef9f6df0 Mon Sep 17 00:00:00 2001
From: Keno Fischer <keno@juliacomputing.com>
Date: Sat, 17 Jun 2023 04:16:27 +0000
Subject: [PATCH] Set u_modified in DAE initialization if the u was changed

Some integrators (such as Sundials) keep a shadow copy of `u` and need
to reflect this copy back into the solver internals to use them.
This happens e.g. during callbacks, but the same situation applies
to initialization, so use the same mechanism (setting u_modified
to `true`) to let the solver know that it needs to update its
internal datastructures.
---
 src/initialize_dae.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/initialize_dae.jl b/src/initialize_dae.jl
index af4f96ceae..9b7d4610d8 100644
--- a/src/initialize_dae.jl
+++ b/src/initialize_dae.jl
@@ -183,6 +183,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
         integrator.u .= nlsol.u
         failed = nlsol.retcode != ReturnCode.Success
     end
+    u_modified!(integrator, true)
     recursivecopy!(integrator.uprev, integrator.u)
     if alg_extrapolates(integrator.alg)
         recursivecopy!(integrator.uprev2, integrator.uprev)
@@ -233,6 +234,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
         olddt
         failed = nlsolvefail(nlsolver)
         @.. broadcast=false integrator.u=integrator.uprev + z
+        u_modified!(integrator, true)
     else
         nlequation_oop = @closure (u, _) -> begin
             update_coefficients!(M, u, p, t)
@@ -258,6 +260,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
         integrator.u = nlsol.u
         failed = nlsol.retcode != ReturnCode.Success
     end
+    u_modified!(integrator, true)
 
     integrator.uprev = copy(integrator.u)
     if alg_extrapolates(integrator.alg)
@@ -320,6 +323,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         reltol = integrator.opts.reltol)
 
     integrator.u = nlsol.u
+    u_modified!(integrator, true)
     recursivecopy!(integrator.uprev, integrator.u)
     if alg_extrapolates(integrator.alg)
         recursivecopy!(integrator.uprev2, integrator.uprev)
@@ -367,6 +371,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         reltol = integrator.opts.reltol)
 
     integrator.u = nlsol.u
+    u_modified!(integrator, true)
 
     integrator.uprev = copy(integrator.u)
     if alg_extrapolates(integrator.alg)
@@ -443,6 +448,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     nlprob = NonlinearProblem(nlfunc, alg_u, p)
     nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol)
     alg_u .= nlsol
+    u_modified!(integrator, true)
 
     recursivecopy!(integrator.uprev, integrator.u)
     if alg_extrapolates(integrator.alg)
@@ -512,6 +518,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     else
         integrator.u = u
     end
+    u_modified!(integrator, true)
 
     integrator.uprev = copy(integrator.u)
     if alg_extrapolates(integrator.alg)
@@ -573,6 +580,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
 
     @. du = ifelse(differential_vars, nlsol.u, du)
     @. u = ifelse(differential_vars, u, nlsol.u)
+    u_modified!(integrator, true)
 
     recursivecopy!(integrator.uprev, integrator.u)
     if alg_extrapolates(integrator.alg)
@@ -630,6 +638,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
         integrator.u = u
         integrator.du = du
     end
+    u_modified!(integrator, true)
 
     integrator.uprev = copy(integrator.u)
     if alg_extrapolates(integrator.alg)