support for re-solving for Initials in remake for nonlinear systems

[isaacsas]

As discussed on Slack, it would be nice if the following cases all worked, the motivation being that a user might want to change the conservation equation parameter, gamma, and have an unknown of their choosing automatically adjust its u0 value to compensate. This is currently not possible as remake on nonlinear systems only handles parameter reinitialization.

Example:

using ModelingToolkit, NonlinearSolve
    @variables X1 X2 X3
    @parameters k1 k2 k3 k4 Γ[1:1] = missing [guess = ones(1)]
    eqs = [
        0 ~ -k1*X1 + k2*X2 - k3*X1*X2 + (1//2)*k4*((-X1 - X2 + Γ[1])^2),
        0 ~ k1*X1 - k2*X2 - k3*X1*X2 + (1//2)*k4*((-X1 - X2 + Γ[1])^2),
        0 ~ -X1 - X2 - X3 + Γ[1]
    ]
    initeqs = [Γ[1] ~ Initial(X1) + Initial(X3) + Initial(X2)]
    @named nlsys = NonlinearSystem(eqs, [X1, X2, X3], [k1, k2, k3, k4, Γ]; 
        initialization_eqs = initeqs)

    u0 = [:X1 => 1.0, :X2 => 2.0, :X3 => 3.0]
    ps = [:k1 => 0.1, :k2 => 0.2, :k3 => 0.3, :k4 => 0.4]
    
    # WITHOUT structural_simplify
    nlsys1 = complete(nlsys)
    nlprob1 = NonlinearProblem(nlsys1, u0, ps)

    # these pass
    @test nlprob1[:X1] == 1.0
    @test nlprob1[:X2] == 2.0
    @test nlprob1[:X3] == 3.0
    @test nlprob1.ps[:Γ][1] == 6.0
    integ1 = init(nlprob1, NewtonRaphson())
    @test integ1[:X1] == 1.0
    @test integ1[:X2] == 2.0
    @test integ1[:X3] == 3.0
    @test integ1.ps[:Γ][1] == 6.0

    nlprob1b = remake(nlprob1; u0 = [:X3 => nothing], p = [:Γ => [4.0]])
    @test nlprob1b[:X1] == 1.0   # pass
    @test nlprob1b[:X2] == 2.0   # pass
    @test_broken nlprob1b[:X3] == 1.0   # fails as still 3.0
    @test nlprob1b.ps[:Γ][1] == 4.0 # pass
    integ1 = init(nlprob1b, NewtonRaphson())
    @test integ1[:X1] == 1.0  # pass
    @test integ1[:X2] == 2.0  # pass
    @test_broken integ1[:X3] == 1.0  # fails
    @test integ1.ps[:Γ][1] == 4.0 # pass

    nlprob1c = remake(nlprob1; u0 = [:X2 => nothing], p = [:Γ => [4.0]])
    @test nlprob1c[:X1] == 1.0   # pass
    @test nlprob1c[:X2] == 0.0   # pass
    @test nlprob1c[:X3] == 3.0   # fails as still 3.0
    @test nlprob1c.ps[:Γ][1] == 4.0 # pass
    integ1 = init(nlprob1b, NewtonRaphson())
    @test integ1[:X1] == 1.0  # pass
    @test integ1[:X2] == 2.0  # pass
    @test_broken integ1[:X3] == 1.0  # fails
    @test integ1.ps[:Γ][1] == 4.0 # pass
    
    # WITH structural_simplify
    nlsys2 = structural_simplify(nlsys)
    nlprob2 = NonlinearProblem(nlsys2, u0, ps)

    # these pass
    @test nlprob2[:X1] == 1.0
    @test nlprob2[:X2] == 2.0
    @test nlprob2[:X3] == 3.0
    @test nlprob2.ps[:Γ][1] == 6.0
    integ2 = init(nlprob2, NewtonRaphson())
    @test integ2[:X1] == 1.0
    @test integ2[:X2] == 2.0
    @test integ2[:X3] == 3.0
    @test integ2.ps[:Γ][1] == 6.0

    nlprob2b = remake(nlprob2; u0 = [:X3 => nothing], p = [:Γ => [4.0]])
    @test nlprob2b[:X1] == 1.0   # pass
    @test nlprob2b[:X2] == 2.0   # pass
    @test nlprob2b[:X3] == 1.0   # pass
    @test nlprob2b.ps[:Γ][1] == 4.0 # pass
    integ2 = init(nlprob2b, NewtonRaphson())
    @test integ2[:X1] == 1.0  # pass
    @test integ2[:X2] == 2.0  # pass
    @test integ2[:X3] == 1.0  # pass
    @test integ2.ps[:Γ][1] == 4.0 # pass

    nlprob2c = remake(nlprob2; u0 = [:X2 => nothing], p = [:Γ => [4.0]])
    @test nlprob2c[:X1] == 1.0   # pass
    @test nlprob2c[:X2] == 0.0   # fails
    @test nlprob2c[:X3] == 3.0   # fails
    @test nlprob2c.ps[:Γ][1] == 4.0 # pass
    integ2 = init(nlprob2c, NewtonRaphson())
    @test integ2[:X1] == 1.0  # pass
    @test integ2[:X2] == 0.0  # fails
    @test integ2[:X3] == 3.0  # fails
    @test integ2.ps[:Γ][1] == 4.0 # pass

On Slack @AayushSabharwal concluded the update should be::

"Okay so if a variable x is missing a u0 , and Initial(x) occurs in an initialization equation we make it solvable despite it not having a guess and let the InitializationProblem constructor error if it ends up in the unknowns?"

with @ChrisRackauckas saying:

"Yes. Just like how with other things in a nonlinear solve, we let it work if they don't have a guess and it simplifies"

and Chris suggesting remake should work like

nlprob1 = remake(nlprob1; u0 = [:X3 => nothing], p = [:Γ => [4.0], Initial(X3) => nothing])

for the desired u0 updates.