Support save_idxs with symbolic indexing

[TorkelE]

Compiling a list of stuff that has been broken since he latest update. Noteworthy is that remake fails in diverse way depending on what problem you are using. Also, the new getp and setp interface functions are a bit weird (see SciML/SymbolicIndexingInterface.jl#24 for more details).

There is a lot of repetitive code in here, however, there are some unusual cases where SDE solutions fail for plotting, so probably make sense to go through it all.

Everything that does not have a comment is working fine. Lines with potential errors or weird behaviours are commented.

using Catalyst, OrdinaryDiffEq, StochasticDiffEq, JumpProcesses, NonlinearSolve, Plots
import ModelingToolkit: getp, setp

model = complete(@reaction_network begin
    (kp,kd), 0 <--> X
    (k1,k2), X <--> Y
end)
@unpack X, Y, kp, kd, k1, k2 = model

u0 = [X => 0.1, Y => 0.0]
tspan = (0.0, 10.0)
ps = [kp => 1.0, kd => 0.2, k1 => 1.0, k2 => 2.0]

# ODEProblem
oprob = ODEProblem(model, u0, tspan, ps)
oprob[X]
oprob[model.X]
oprob[:X]
oprob[[X,Y]]
oprob[[model.X,model.Y]]
oprob[[:X,:Y]]
oprob[(X,Y)]                # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
oprob[(model.X,model.Y)]    # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
oprob[(:X,:Y)]              # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.

oprob[X] = 0.2
oprob[model.X] = 0.3
oprob[:X] = 0.4

# Yes, here and in alter cases the first 3 have the intended behaviour, but still feels like a really weird choice of function name, see previous comment.
getp(oprob, kp) # Returns a SymbolicIndexingInterface struct.
getp(oprob, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(oprob, :kp) # Returns a SymbolicIndexingInterface struct.
getp(oprob, kp)(oprob)
getp(oprob, model.kp)(oprob)
getp(oprob, :kp)(oprob)

setp(oprob, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(oprob, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(oprob, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(oprob, kp)(oprob, 5.0)
setp(oprob, model.kp)(oprob, 6.0)
setp(oprob, :kp)(oprob, 7.0)

remake(oprob; u0 = [X => 0.2, Y => 0.2]).u0
remake(oprob; u0 = [model.X => 0.3, model.Y => 0.3]).u0
remake(oprob; u0 = [:X => 0.4, :Y => 0.4]).u0 # Does not update properly.
remake(oprob; u0 = [X => 0.5]).u0
remake(oprob; u0 = [model.X => 0.6]).u0
remake(oprob; u0 = [:X => 0.7]).u0 # Does not update properly.

remake(oprob; p = [kp => 0.1, kd => 0.1, k1 => 0.1, k2 => 0.1]).p
remake(oprob; p = [model.kp => 0.2, model.kd => 0.2, model.k1 => 0.2, model.k2 => 0.2]).p
remake(oprob; p = [:kp => 0.3, :kd => 0.3, :k1 => 0.3, :k2 => 0.3]).p # ERROR: ArgumentError: Any[kd, k1, k2] are missing from the variable map.
remake(oprob; p = [kp => 0.4]).p # ERROR: ArgumentError: SymbolicUtils.BasicSymbolic{Real}[kd, k1, k2] are missing from the variable map.
remake(oprob; p = [model.kp => 0.5]).p # ERROR: ArgumentError: SymbolicUtils.BasicSymbolic{Real}[kd, k1, k2] are missing from the variable map.
remake(oprob; p = [:kp => 0.6]).p # ERROR: ArgumentError: Any[kd, k1, k2] are missing from the variable map.

# SDEProblem
sprob = SDEProblem(model, u0, tspan, ps)
sprob[X]
sprob[model.X]
sprob[:X]
sprob[[X,Y]]
sprob[[model.X,model.Y]]
sprob[[:X,:Y]]
sprob[(X,Y)]                # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
sprob[(model.X,model.Y)]    # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
sprob[(:X,:Y)]              # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.

sprob[X] = 0.2
sprob[model.X] = 0.3
sprob[:X] = 0.4

getp(sprob, kp) # Returns a SymbolicIndexingInterface struct.
getp(sprob, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(sprob, :kp) # Returns a SymbolicIndexingInterface struct.
getp(sprob, kp)(sprob)
getp(sprob, model.kp)(sprob)
getp(sprob, :kp)(sprob)

setp(sprob, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(sprob, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(sprob, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(sprob, kp)(sprob, 5.0)
setp(sprob, model.kp)(sprob, 6.0)
setp(sprob, :kp)(sprob, 7.0)

remake(sprob; u0 = [X => 0.2, Y => 0.2]).u0 # Makes u0 a vector of pairs (should be a normal vector).
remake(sprob; u0 = [model.X => 0.3, model.Y => 0.3]).u0 # Makes u0 a vector of pairs (should be a normal vector).
remake(sprob; u0 = [:X => 0.4, :Y => 0.4]).u0 # Makes u0 a vector of pairs (should be a normal vector).
remake(sprob; u0 = [X => 0.5]).u0 # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(sprob; u0 = [model.X => 0.6]).u0 # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(sprob; u0 = [:X => 0.7]).u0 # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.

remake(sprob; p = [kp => 0.1, kd => 0.1, k1 => 0.1, k2 => 0.1]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(sprob; p = [model.kp => 0.2, model.kd => 0.2, model.k1 => 0.2, model.k2 => 0.2]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(sprob; p = [:kp => 0.3, :kd => 0.3, :k1 => 0.3, :k2 => 0.3]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(sprob; p = [kp => 0.4]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(sprob; p = [model.kp => 0.5]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(sprob; p = [:kp => 0.6]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.

# DiscreteProblem
dprob = DiscreteProblem(model, [X => 1, Y => 0], tspan, ps)
dprob[X]
dprob[model.X]
dprob[:X]
dprob[[X,Y]]
dprob[[model.X,model.Y]]
dprob[[:X,:Y]]
dprob[(X,Y)]                # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
dprob[(model.X,model.Y)]    # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
dprob[(:X,:Y)]              # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.

dprob[X] = 2
dprob[model.X] = 3
dprob[:X] = 4

getp(dprob, kp) # Returns a SymbolicIndexingInterface struct.
getp(dprob, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(dprob, :kp) # Returns a SymbolicIndexingInterface struct.
getp(dprob, kp)(dprob)
getp(dprob, model.kp)(dprob)
getp(dprob, :kp)(dprob)

setp(dprob, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(dprob, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(dprob, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(dprob, kp)(sprob, 5.0)
setp(dprob, model.kp)(sprob, 6.0)
setp(dprob, :kp)(sprob, 7.0)

remake(dprob; u0 = [X => 0.2, Y => 0.2]).u0 # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; u0 = [model.X => 0.3, model.Y => 0.3]).u0 # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; u0 = [:X => 0.4, :Y => 0.4]).u0 # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; u0 = [X => 0.5]).u0 # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(dprob; u0 = [model.X => 0.6]).u0 # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(dprob; u0 = [:X => 0.7]).u0 # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.

remake(dprob; p = [kp => 0.1, kd => 0.1, k1 => 0.1, k2 => 0.1]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; p = [model.kp => 0.2, model.kd => 0.2, model.k1 => 0.2, model.k2 => 0.2]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; p = [:kp => 0.3, :kd => 0.3, :k1 => 0.3, :k2 => 0.3]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; p = [kp => 0.4]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(dprob; p = [model.kp => 0.5]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(dprob; p = [:kp => 0.6]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.

# JumpProblem
jprob = JumpProblem(model, dprob, Direct())
jprob[X]
jprob[model.X]
jprob[:X]
jprob[[X,Y]]
jprob[[model.X,model.Y]]
jprob[[:X,:Y]]
jprob[(X,Y)]                # Error. Not really needed, but supported for solutions.
jprob[(model.X,model.Y)]    # Error. Not really needed, but supported for solutions.
jprob[(:X,:Y)]              # Error. Not really needed, but supported for solutions.

jprob[X] = 2
jprob[model.X] = 3
jprob[:X] = 4

getp(jprob, kp) # ERROR: type JumpProblem has no field f
getp(jprob, model.kp) # ERROR: type JumpProblem has no field f
getp(jprob, :kp) # ERROR: type JumpProblem has no field f
getp(jprob, kp)(jprob) # ERROR: type JumpProblem has no field f
getp(jprob, model.kp)(jprob) # ERROR: type JumpProblem has no field f
getp(jprob, :kp)(jprob) # ERROR: type JumpProblem has no field f

setp(jprob, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(jprob, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(jprob, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(jprob, kp)(jprob, 5.0) # ERROR: type JumpProblem has no field f
setp(jprob, model.kp)(jprob, 6.0) # ERROR: type JumpProblem has no field f
setp(jprob, :kp)(jprob, 7.0) # ERROR: type JumpProblem has no field f

remake(jprob; u0 = [X => 0.2, Y => 0.2]).u0 # ERROR: type JumpProblem has no field u0
remake(jprob; u0 = [model.X => 0.3, model.Y => 0.3]).u0 # ERROR: type JumpProblem has no field u0
remake(jprob; u0 = [:X => 0.4, :Y => 0.4]).u0 # ERROR: type JumpProblem has no field u0
remake(jprob; u0 = [X => 0.5]).u0 # ERROR: type JumpProblem has no field u0
remake(jprob; u0 = [model.X => 0.6]).u0 # ERROR: type JumpProblem has no field u0
remake(jprob; u0 = [:X => 0.7]).u0 # ERROR: type JumpProblem has no field u0

remake(dprob; p = [kp => 0.1, kd => 0.1, k1 => 0.1, k2 => 0.1]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; p = [model.kp => 0.2, model.kd => 0.2, model.k1 => 0.2, model.k2 => 0.2]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; p = [:kp => 0.3, :kd => 0.3, :k1 => 0.3, :k2 => 0.3]).p # Makes u0 a vector of pairs (should be a normal vector).
remake(dprob; p = [kp => 0.4]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(dprob; p = [model.kp => 0.5]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.
remake(dprob; p = [:kp => 0.6]).p # Makes u0 a vector of pairs (should be a normal vector). Also makes it a vector with value for X only.

# NonlinearProblem
nlprob = NonlinearProblem(model, u0, ps)
nlprob[X]
nlprob[model.X]
nlprob[:X]
nlprob[[X,Y]]
nlprob[[model.X,model.Y]]
nlprob[[:X,:Y]]
nlprob[(X,Y)]                # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
nlprob[(model.X,model.Y)]    # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.
nlprob[(:X,:Y)]              # ERROR: Invalid indexing of problem. Not really needed, but supported for solutions and integrators.

nlprob[X] = 0.2
nlprob[model.X] = 0.3
nlprob[:X] = 0.4

getp(nlprob, kp) # Returns a SymbolicIndexingInterface struct.
getp(nlprob, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(nlprob, :kp) # Returns a SymbolicIndexingInterface struct.
getp(nlprob, kp)(nlprob)
getp(nlprob, model.kp)(nlprob)
getp(nlprob, :kp)(nlprob)

setp(nlprob, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(nlprob, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(nlprob, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(nlprob, kp)(nlprob, 5.0)
setp(nlprob, model.kp)(nlprob, 6.0)
setp(nlprob, :kp)(nlprob, 7.0)

remake(nlprob; u0 = [X => 0.2, Y => 0.2]).u0
remake(nlprob; u0 = [model.X => 0.3, model.Y => 0.3]).u0
remake(nlprob; u0 = [:X => 0.4, :Y => 0.4]).u0 # Does not update properly.
remake(nlprob; u0 = [X => 0.5]).u0
remake(nlprob; u0 = [model.X => 0.6]).u0
remake(nlprob; u0 = [:X => 0.7]).u0 # Does not update properly.

remake(nlprob; p = [kp => 0.1, kd => 0.1, k1 => 0.1, k2 => 0.1]).p
remake(nlprob; p = [model.kp => 0.2, model.kd => 0.2, model.k1 => 0.2, model.k2 => 0.2]).p
remake(nlprob; p = [:kp => 0.3, :kd => 0.3, :k1 => 0.3, :k2 => 0.3]).p  # Does not update properly.
remake(nlprob; p = [kp => 0.4]).p
remake(nlprob; p = [model.kp => 0.5]).p
remake(nlprob; p = [:kp => 0.6]).p  # Does not update properly.


# Integrator
using DifferentialEquations
integrator = init(oprob)

ointegrator = init(ODEProblem(model, u0, tspan, ps))
ointegrator[X]
ointegrator[model.X]
ointegrator[:X]
ointegrator[[X,Y]]
ointegrator[[model.X,model.Y]]
ointegrator[[:X,:Y]]
ointegrator[(X,Y)]              
ointegrator[(model.X,model.Y)] 
ointegrator[(:X,:Y)] 

ointegrator[X] = 0.2
ointegrator[model.X] = 0.3
ointegrator[:X] = 0.4

getp(ointegrator, kp) # Returns a SymbolicIndexingInterface struct.
getp(ointegrator, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(ointegrator, :kp) # Returns a SymbolicIndexingInterface struct.
getp(ointegrator, kp)(ointegrator)
getp(ointegrator, model.kp)(ointegrator)
getp(ointegrator, :kp)(ointegrator)

setp(ointegrator, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(ointegrator, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(ointegrator, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(ointegrator, kp)(ointegrator, 2.0)
setp(ointegrator, model.kp)(ointegrator, 3.0)
setp(ointegrator, :kp)(ointegrator, 4.0)

sintegrator = init(SDEProblem(model, u0, tspan, ps))
sintegrator[X]
sintegrator[model.X]
sintegrator[:X]
sintegrator[[X,Y]]
sintegrator[[model.X,model.Y]]
sintegrator[[:X,:Y]]
sintegrator[(X,Y)]              
sintegrator[(model.X,model.Y)] 
sintegrator[(:X,:Y)] 

sintegrator[X] = 0.2
sintegrator[model.X] = 0.3
sintegrator[:X] = 0.4

getp(sintegrator, kp) # Returns a SymbolicIndexingInterface struct.
getp(sintegrator, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(sintegrator, :kp) # Returns a SymbolicIndexingInterface struct.
getp(sintegrator, kp)(sintegrator)
getp(sintegrator, model.kp)(sintegrator)
getp(sintegrator, :kp)(sintegrator)

setp(sintegrator, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(sintegrator, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(sintegrator, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(sintegrator, kp)(sintegrator, 2.0)
setp(sintegrator, model.kp)(sintegrator, 3.0)
setp(sintegrator, :kp)(sintegrator, 4.0)

jintegrator = init(jprob, SSAStepper())
jintegrator[X]
jintegrator[model.X]
jintegrator[:X]
jintegrator[[X,Y]]
jintegrator[[model.X,model.Y]]
jintegrator[[:X,:Y]]
jintegrator[(X,Y)]              
jintegrator[(model.X,model.Y)] 
jintegrator[(:X,:Y)] 

jintegrator[X] = 2
jintegrator[model.X] = 3
jintegrator[:X] = 4

getp(jintegrator, kp) # Returns a SymbolicIndexingInterface struct.
getp(jintegrator, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(jintegrator, :kp) # Returns a SymbolicIndexingInterface struct.
getp(jintegrator, kp)(jintegrator)
getp(jintegrator, model.kp)(jintegrator)
getp(jintegrator, :kp)(jintegrator)

setp(jintegrator, kp, 2.0) # ERROR: MethodError: no method matching setp(...
setp(jintegrator, model.kp, 3.0) # ERROR: MethodError: no method matching setp(...
setp(jintegrator, :kp, 4.0) # ERROR: MethodError: no method matching setp(...
setp(jintegrator, kp)(jintegrator, 2.0)
setp(jintegrator, model.kp)(jintegrator, 3.0)
setp(jintegrator, :kp)(jintegrator, 4.0)

# Solve command
solve(oprob, Tsit5(); save_idxs=[X]) # ERROR: TypeError: non-boolean (Num) used in boolean context
solve(sprob, ImplicitEM(); save_idxs=[X]) # ERROR: TypeError: non-boolean (Num) used in boolean context
solve(jprob, SSAStepper(); save_idxs=[X]) # ERROR: MethodError: no method matching __init(::JumpProblem{...

solve(oprob, Tsit5(); save_idxs=[model.X]) # ERROR: TypeError: non-boolean (Num) used in boolean context
solve(sprob, ImplicitEM(); save_idxs=[model.X]) # ERROR: TypeError: non-boolean (Num) used in boolean context
solve(jprob, SSAStepper(); save_idxs=[model.X]) # ERROR: MethodError: no method matching __init(::JumpProblem{

solve(oprob, Tsit5(); save_idxs=[:X]) # ERROR: ArgumentError: unable to check bounds for indices of type Symbol
solve(sprob, ImplicitEM(); save_idxs=[:X]) # ERROR: ArgumentError: unable to check bounds for indices of type Symbol
solve(jprob, SSAStepper(); save_idxs=[:X]) # ERROR: MethodError: no method matching __init(::JumpProblem{true, 

# Fore reference, normal indexing also fails here:
solve(jprob, SSAStepper(); save_idxs=[1]) # ERROR: MethodError: no method matching __init(::JumpProblem{true, 

# Solution
osol = solve(oprob, Tsit5())
ssol = solve(sprob, ImplicitEM())
jsol = solve(jprob, SSAStepper())
nlsol = solve(nlprob, NewtonRaphson())

osol[X]
osol[model.X]
osol[:X]
osol[[X,Y]]
osol[[model.X,model.Y]]
osol[[:X,:Y]]
osol[(X,Y)]
osol[(model.X,model.Y)]
osol[(:X,:Y)]

getp(osol, kp) # Returns a SymbolicIndexingInterface struct.
getp(osol, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(osol, :kp) # Returns a SymbolicIndexingInterface struct.
getp(osol, kp)(osol)
getp(osol, model.kp)(osol)
getp(osol, :kp)(osol)

ssol[X]
ssol[model.X]
ssol[:X]
ssol[[X,Y]]
ssol[[model.X,model.Y]]
ssol[[:X,:Y]]
ssol[(X,Y)]
ssol[(model.X,model.Y)]
ssol[(:X,:Y)]

getp(ssol, kp) # Returns a SymbolicIndexingInterface struct.
getp(ssol, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(ssol, :kp) # Returns a SymbolicIndexingInterface struct.
getp(ssol, kp)(osol)
getp(ssol, model.kp)(osol)
getp(ssol, :kp)(osol)

jsol[X]
jsol[model.X]
jsol[:X]
jsol[[X,Y]]
jsol[[model.X,model.Y]]
jsol[[:X,:Y]]
jsol[(X,Y)]
jsol[(model.X,model.Y)]
jsol[(:X,:Y)]

getp(jsol, kp) # Returns a SymbolicIndexingInterface struct.
getp(jsol, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(jsol, :kp) # Returns a SymbolicIndexingInterface struct.
getp(jsol, kp)(osol)
getp(jsol, model.kp)(osol)
getp(jsol, :kp)(osol)

nlsol[X]
nlsol[model.X]
nlsol[:X]
nlsol[[X,Y]]
nlsol[[model.X,model.Y]]
nlsol[[:X,:Y]]
nlsol[(X,Y)] # ERROR: Invalid indexing of solution
nlsol[(model.X,model.Y)] # ERROR: Invalid indexing of solution
nlsol[(:X,:Y)] # ERROR: Invalid indexing of solution

getp(nlsol, kp) # Returns a SymbolicIndexingInterface struct.
getp(nlsol, model.kp) # Returns a SymbolicIndexingInterface struct.
getp(nlsol, :kp) # Returns a SymbolicIndexingInterface struct.
getp(nlsol, kp)(osol)
getp(nlsol, model.kp)(osol)
getp(nlsol, :kp)(osol)


# Plot command
plot(osol; idxs = [X])
plot(osol; idxs = [model.X])
plot(osol; idxs = [:X])

plot(ssol; idxs = [X]) # ERROR: ArgumentError: invalid index: X(t) of type 
plot(ssol; idxs = [model.X]) # ERROR: ArgumentError: invalid index: X(t) of type 
plot(ssol; idxs = [:X]) # ERROR: ArgumentError: invalid index: :X of type Symbol

plot(jsol; idxs = [X])
plot(jsol; idxs = [model.X])
plot(jsol; idxs = [:X])

[TorkelE]

For reference, here is a full test set of various indexing features that I would like to run once we are done with everything. It is using Catalyst currently (but I was unsure how to create symbolic-type JumpProblems without it).

# Fetch packages
using Catalyst, OrdinaryDiffEq, StochasticDiffEq, JumpProcesses, NonlinearSolve, Plots
using Test

# Create model, problems, and solutions.
begin
    model = complete(@reaction_network begin
        (kp,kd), 0 <--> X
        (k1,k2), X <--> Y
    end)
    @unpack X, Y, kp, kd, k1, k2 = model

    u0_vals = [X => 1, Y => 0]
    tspan = (0.0, 10.0)
    p_vals = [kp => 1.0, kd => 0.2, k1 => 1.0, k2 => 2.0]

    oprob = ODEProblem(model, u0_vals, tspan, p_vals)
    sprob = SDEProblem(model,u0_vals, tspan, p_vals)
    dprob = DiscreteProblem(model, u0_vals, tspan, p_vals)
    jprob = JumpProblem(model, dprob, Direct())
    nprob = NonlinearProblem(model, u0_vals, p_vals)
    problems = [oprob, sprob, dprob, jprob, nprob]
    
    osol = solve(oprob, Tsit5())
    ssol = solve(sprob, ImplicitEM())
    jsol = solve(jprob, SSAStepper())
    nsol = solve(nprob, NewtonRaphson())
    sols = [osol, ssol, jsol, nsol]
end

# Tests index updating.
let 
    for prob in deepcopy(problems)
        # Get u values.
        @test prob[X] == prob[model.X] == prob[:X] == 1
        @test prob[[X,Y]] == prob[[model.X,model.Y]] == prob[[:X,:Y]] == [1, 2]
        @test prob[(X,Y)] == prob[(model.X,model.Y)] == prob[(:X,:Y)] == (1, 2)
        @test getu(prob, X)(prob) == getu(prob, model.X)(prob) == getu(prob, :X)(prob) == 1      

        # Set u values.
        prob[X] = 2
        @test prob[X] == 2
        prob[model.X] = 3
        @test prob[X] == 3
        prob[:X] = 4
        @test prob[X] == 4
        setu(prob, X)(prob, 5)
        @test prob[X] == 5
        setu(prob, model.X)(prob, 6)
        @test prob[X] == 6
        setu(prob, :X)(prob, 7)
        @test prob[X] == 7

        # Get p values.
        @test getp(prob, kp)(prob) == getp(prob, model.kp)(prob) == getp(prob, :kp)(prob) == 1.0
        @test prob.ps[kp] == prob.ps[model.kp] == prob.ps[:kp] == 1.0    
        
        # Set p values.
        setp(prob, kp)(prob, 2.0)
        @test prob[kp] == 2.0
        setp(prob, model.kp)(prob, 3.0)
        @test prob[kp] == 3.0
        setp(prob, :kp)(prob, 4.0)
        @test prob[kp] == 4.0
        prob.ps[kp] = 5.0
        @test prob[kp] == .0
        prob.ps[model.kp] = 6.0
        @test prob[kp] == 6.0
        prob.ps[:kp] = 7.0
        @test prob[kp] == 7.0
    end
end

# Test remake function.
let 
    for prob in deepcopy(probs)
        # Remake for all u0s.
        @test remake(prob; u0 = [X => 2, Y => 3]).u0 == [2, 3]
        @test remake(prob; u0 = [model.X => 4, model.Y => 5]).u0 == [4, 5]
        @test remake(prob; u0 = [:X => 6, :Y => 7]).u0 == [6, 7]

        # Remake for some u0s.
        @test remake(prob; u0 = [Y => 8]).u0 == [1, 8]
        @test remake(prob; u0 = [model.Y => 9]).u0 == [1, 9]
        @test remake(prob; u0 = [:Y => 10]).u0 == [1, 10]

        # Remake for all ps.
        @test remake(prob; p = [kp => 1.0, kd => 2.0, k1 => 3.0, k2 => 4.0]).p == [1.0, 2.0, 3.0, 4.0]
        @test remake(prob; p = [model.kp => 5.0, model.kd => 6.0, model.k1 => 7.0, model.k2 => 8.0]).p == [5.0, 6.0, 7.0, 8.0]
        @test remake(prob; p = [:kp => 9.0, :kd => 10.0, :k1 => 11.0, :k2 => 12.0]).p  == [9.0, 10.0, 11.0, 12.0]

        # Remake for some ps.
        @test remake(prob; p = [k2 => 13.0]).p == [1.0, 0.2, 1.0, 13.0]
        @test remake(prob; p = [model.k2 => 14.0]).p == [1.0, 0.2, 1.0, 14.0]
        @test remake(prob; p = [:k2 => 15.0]).p  == [1.0, 0.2, 1.0, 15.0]
    end
end

# Test integrator indexing.
let 
    for integrator in init.(deepcopy(problems))
        # Get u values.
        @test integrator[X] == integrator[model.X] == integrator[:X] == 1
        @test integrator[[X,Y]] == integrator[[model.X,model.Y]] == integrator[[:X,:Y]] == [1, 2]
        @test integrator[(X,Y)] == integrator[(model.X,model.Y)] == integrator[(:X,:Y)] == (1, 2)   
        @test getu(integrator, X)(integrator) == getu(integrator, model.X)(integrator) == getu(integrator, :X)(integrator) == 1         

        # Set u values.
        integrator[X] = 2
        @test integrator[X] == 2
        integrator[model.X] = 3
        @test integrator[X] == 3
        integrator[:X] = 4
        @test integrator[X] == 4
        setu(integrator, X)(integrator, 5)
        @test integrator[X] == 5
        setu(integrator, model.X)(integrator, 6)
        @test integrator[X] == 6
        setu(integrator, :X)(integrator, 7)
        @test integrator[X] == 7

        # Get p values.
        @test getp(integrator, kp)(integrator) == getp(integrator, model.kp)(integrator) == getp(integrator, :kp)(integrator) == 1.0
        @test integrator.ps[kp] == integrator.ps[model.kp] == integrator.ps[:kp] == 1.0    

        # Set p values.
        setp(integrator, kp)(integrator, 2.0)
        @test integrator[kp] == 2.0
        setp(integrator, model.kp)(integrator, 3.0)
        @test integrator[kp] == 3.0
        setp(integrator, :kp)(integrator, 4.0)integrator
        @test integrator[kp] == 4.0
        integrator.ps[kp] = 5.0
        @test integrator[kp] == .0
        integrator.ps[model.kp] = 6.0
        @test integrator[kp] == 6.0
        integrator.ps[:kp] = 7.0
        @test integrator[kp] == 7.0
    end
end

# Test solve's save_idxs argument.
let 
    @test length(solve(oprob, Tsit5(); save_idxs=[X]).u[1]) == 1
    @test length(solve(sprob, ImplicitEM(); save_idxs=[X]).u[1]) == 1
    @test length(solve(jprob, SSAStepper(); save_idxs=[X]).u[1]) == 1
end

#Tests solution indexing.
let 
    for sol in deepcopy(sols)
        # Get u values.
        @test Int64(sol[X][1]) == 1
        @test Int64(sol[Y][1]) == 0
        @test sol[X] == sol[model.X] == sol[:X]
        @test sol[[X,Y]] == sol[[model.X,model.Y]] == sol[[:X,:Y]]
        @test sol[(X,Y)] == sol[(model.X,model.Y)] == sol[(:X,:Y)]   
        @test getu(sol, X)(sol) == getu(sol, model.X)(sol) == getu(sol, :X)(sol)          

        # Get p values.
        @test getp(sol, kp)(sol) == getp(sol, model.kp)(sol) == getp(sol, :kp)(sol) == 1.0
        @test sol.ps[kp] == sol.ps[model.kp] == sol.ps[:kp] == 1.0    
    end
end

# Tests plotting.
let 
    for sol in deepcopy(sols[1:3])
        # Single variable.
        @test length(plot(sol; idxs = X).series_list) == 1
        @test length(plot(sol; idxs = model.X).series_list) == 1
        @test length(plot(sol; idxs = :X).series_list) == 1

        # As vector.
        @test length(plot(sol; idxs = [X,Y]).series_list) == 2
        @test length(plot(sol; idxs = [model.X,model.Y]).series_list) == 2
        @test length(plot(sol; idxs = [:X,:Y]).series_list) == 2

        # As tuple.
        @test length(plot(sol; idxs = (X, Y)).series_list) == 1
        @test length(plot(sol; idxs = (model.X, model.Y)).series_list) == 1
        @test length(plot(sol; idxs = (:X, :Y)).series_list) == 1
    end    
end

# Tests solving for various inputs types.
let 
    u0_vals = [X => 1, Y => 0]
    tspan = (0.0, 10.0)
    p_vals = [kp => 1.0, kd => 0.2, k1 => 1.0, k2 => 2.0]

    u0_vals_2 = [model.X => 1, model.Y => 0]
    u0_vals_3 = [:X => 1, :Y => 0]
    p_vals_2 = [model.kp => 1.0, model.kd => 0.2, model.k1 => 1.0, model.k2 => 2.0]
    p_vals_3 = [:kp => 1.0, :kd => 0.2, :k1 => 1.0, :k2 => 2.0]

    oprob_2 = ODEProblem(model, u0_vals_2, tspan, p_vals_2)
    oprob_3 = ODEProblem(model, u0_vals_3, tspan, p_vals_3)
    sprob_2 = SDEProblem(model,u0_vals_2, tspan, p_vals_2)
    sprob_3 = SDEProblem(model,u0_vals_3, tspan, p_vals_3)
    dprob_2 = DiscreteProblem(model, u0_vals_2, tspan, p_vals_2)
    dprob_3 = DiscreteProblem(model, u0_vals_3, tspan, p_vals_3)
    jprob_2 = JumpProblem(model, dprob_2, Direct())
    jprob_3 = JumpProblem(model, dprob_3, Direct())
    nprob_2 = NonlinearProblem(model, u0_vals_2, p_vals_2)
    nprob_3 = NonlinearProblem(model, u0_vals_3, p_vals_3)
    
    @test solve(oprob, Tsit5()) == solve(oprob_2, Tsit5()) == solve(oprob_3, Tsit5())
    @test solve(sprob, ImplicitEM(); seed=1234) == solve(sprob_2, ImplicitEM(); seed=1234) == solve(sprob_3, ImplicitEM(); seed=1234)
    @test solve(jprob, SSAStepper(); seed=1234) == solve(jprob_2, SSAStepper(); seed=1234) == solve(jprob_3, SSAStepper(); seed=1234)
    @test solve(nprob, NewtonRaphson()) == solve(nprob_2, NewtonRaphson()) == solve(nprob_3, NewtonRaphson())
end

[BambOoxX]

Hi all, I am not entirely sure if this is related, but I get the following error while indexing in a SteadyStateSolution, that I do not get when indexing into a ODESolution.

ERROR: MethodError: no method matching (::SciMLBase.var"#187#189"{ODEFunction{…}, Num})(::Vector{Float64}, ::Vector{Float64})

Closest candidates are:
  (::SciMLBase.var"#187#189")(::Any, ::Any, ::Any)
   @ SciMLBase D:\user\.julia\packages\SciMLBase\dpafx\src\scimlfunctions.jl:4225

Stacktrace:
 [1] getindex(A::SciMLBase.NonlinearSolution{…}, sym::Num)
   @ SciMLBase D:\user\.julia\packages\SciMLBase\dpafx\src\solutions\solution_interface.jl:67
 [2] top-level scope
   @ REPL[23]:1
Some type information was truncated. Use `show(err)` to see complete types.

This happens with

  [0bca4576] SciMLBase v2.17.1
  [961ee093] ModelingToolkit v8.75.0

I can open a separate issue if you think this is not related.

[AayushSabharwal]

Could you provide an MWE @BambOoxX ?

[BambOoxX]

I've tried to, but this error does not show on the simplified systems I have tested so far. It looks like this happens when nesting components in ModelingToolkit. I'll try to reproduce with an example I used on the discourse.

[TorkelE]

Think a decent number of these have been fixed. idxs for SDEProblems is still broken though.

[TorkelE]

Here's an updated list. Everything that does not currently work is marked using @test_broken. I currently use Caalyst to generate the models (as it makes it easy to generate all Problem types in a single declaration). For SciML tests you'd probably want to use MTK to create the various models though.

Things mostly work, however, there are still some cases that does not.

# Fetch packages
using Catalyst, OrdinaryDiffEq, StochasticDiffEq, JumpProcesses, NonlinearSolve, Plots
using ModelingToolkit: getu, setu, getp, setp
using Test

# Create model, problems, and solutions.
begin
    model = complete(@reaction_network begin
        @observables XY ~ X + Y
        (kp,kd), 0 <--> X
        (k1,k2), X <--> Y
    end)
    @unpack XY, X, Y, kp, kd, k1, k2 = model

    u0_vals = [X => 4, Y => 5]
    tspan = (0.0, 10.0)
    p_vals = [kp => 1.0, kd => 0.1, k1 => 0.25, k2 => 0.5]

    oprob = ODEProblem(model, u0_vals, tspan, p_vals)
    sprob = SDEProblem(model,u0_vals, tspan, p_vals)
    dprob = DiscreteProblem(model, u0_vals, tspan, p_vals)
    jprob = JumpProblem(model, deepcopy(dprob), Direct())
    nprob = NonlinearProblem(model, u0_vals, p_vals)
    problems = [oprob, sprob, dprob, jprob, nprob]

    oint = init(oprob, Tsit5(); save_everystep=false)
    sint = init(sprob, ImplicitEM(); save_everystep=false)
    jint = init(jprob, SSAStepper())
    nint = init(nprob, NewtonRaphson(); save_everystep=false)
    integrators = [oint, sint, jint, nint]
    
    osol = solve(oprob, Tsit5())
    ssol = solve(sprob, ImplicitEM())
    jsol = solve(jprob, SSAStepper())
    nsol = solve(nprob, NewtonRaphson())
    sols = [osol, ssol, jsol, nsol]
end

# Tests problem index updating.
let 
    for prob in deepcopy(problems)
        # Get u values (including observables).
        @test prob[X] == prob[model.X] == prob[:X] == 4
        @test prob[XY] == prob[model.XY] == prob[:XY] == 9
        @test prob[[XY,Y]] == prob[[model.XY,model.Y]] == prob[[:XY,:Y]] == [9, 5]
        @test_broken prob[(XY,Y)] == prob[(model.XY,model.Y)] == prob[(:XY,:Y)] == (9, 5)
        @test getu(prob, X)(prob) == getu(prob, model.X)(prob) == getu(prob, :X)(prob) == 4
        @test getu(prob, XY)(prob) == getu(prob, model.XY)(prob) == getu(prob, :XY)(prob) == 9 
        @test getu(prob, [XY,Y])(prob) == getu(prob, [model.XY,model.Y])(prob) == getu(prob, [:XY,:Y])(prob) == [9, 5]  
        @test getu(prob, (XY,Y))(prob) == getu(prob, (model.XY,model.Y))(prob) == getu(prob, (:XY,:Y))(prob) == (9, 5)

        # Set u values.
        prob[X] = 20
        @test prob[X] == 20
        prob[model.X] = 30
        @test prob[X] == 30
        prob[:X] = 40
        @test prob[X] == 40
        setu(prob, X)(prob, 50)
        @test prob[X] == 50
        setu(prob, model.X)(prob, 60)
        @test prob[X] == 60
        setu(prob, :X)(prob, 70)
        @test prob[X] == 70

        # Get p values.
        @test prob.ps[kp] == prob.ps[model.kp] == prob.ps[:kp] == 1.0    
        @test prob.ps[[k1,k2]] == prob.ps[[model.k1,model.k2]] == prob.ps[[:k1,:k2]] == [0.25, 0.5]
        @test prob.ps[(k1,k2)] == prob.ps[(model.k1,model.k2)] == prob.ps[(:k1,:k2)] == (0.25, 0.5)
        @test getp(prob, kp)(prob) == getp(prob, model.kp)(prob) == getp(prob, :kp)(prob) == 1.0
        @test getp(prob, [k1,k2])(prob) == getp(prob, [model.k1,model.k2])(prob) == getp(prob, [:k1,:k2])(prob) == [0.25, 0.5]
        @test getp(prob, (k1,k2))(prob) == getp(prob, (model.k1,model.k2))(prob) == getp(prob, (:k1,:k2))(prob) == (0.25, 0.5)
        
        # Set p values.
        prob.ps[kp] = 2.0
        @test prob.ps[kp] == 2.0
        prob.ps[model.kp] = 3.0
        @test prob.ps[kp] == 3.0
        prob.ps[:kp] = 4.0
        @test prob.ps[kp] == 4.0
        setp(prob, kp)(prob, 5.0)
        @test prob.ps[kp] == 5.0
        setp(prob, model.kp)(prob, 6.0)
        @test prob.ps[kp] == 6.0
        setp(prob, :kp)(prob, 7.0)
        @test prob.ps[kp] == 7.0
    end
end

# Test remake function.
let 
    for prob in deepcopy(problems)
        # Remake for all u0s. Have to hanlde JumpProblem's separately.
        rp = remake(prob; u0 = [X => 1, Y => 2])
        @test (rp isa JumpProblem ? rp.prob : rp).u0 == [1, 2] 
        rp = remake(prob; u0 = [model.X => 3, model.Y => 4])
        @test (rp isa JumpProblem ? rp.prob : rp).u0 == [3, 4]
        rp = remake(prob; u0 = [:X => 5, :Y => 6])
        @test (rp isa JumpProblem ? rp.prob : rp).u0 == [5, 6]

        # Remake for some u0s.
        rp = remake(prob; u0 = [Y => 7])
        @test (rp isa JumpProblem ? rp.prob : rp).u0 == [4, 7]
        rp = remake(prob; u0 = [model.Y => 8])
        @test (rp isa JumpProblem ? rp.prob : rp).u0 == [4, 8]
        rp = remake(prob; u0 = [:Y => 9])
        @test (rp isa JumpProblem ? rp.prob : rp).u0 == [4, 9]

        # Remake for all ps. Have to check each parameter separately, as prob.p is a SciMLStructure-thing.
        rp = remake(prob; p = [kp => 1.0, kd => 2.0, k1 => 3.0, k2 => 4.0])
        @test (rp.ps[kp] == 1.0) && (rp.ps[kd] == 2.0) && (rp.ps[k1] == 3.0) && (rp.ps[k2] == 4.0)
        rp = remake(prob; p = [model.kp => 5.0, model.kd => 6.0, model.k1 => 7.0, model.k2 => 8.0])
        @test (rp.ps[kp] == 5.0) && (rp.ps[kd] == 6.0) && (rp.ps[k1] == 7.0) && (rp.ps[k2] == 8.0)
        rp = remake(prob; p = [:kp => 9.0, :kd => 10.0, :k1 => 11.0, :k2 => 12.0])
        @test (rp.ps[kp] == 9.0) && (rp.ps[kd] == 10.0) && (rp.ps[k1] == 11.0) && (rp.ps[k2] == 12.0)

        # Remake for some ps.
        rp = remake(prob; p = [k2 => 13.0])
        @test (rp.ps[kp] == 1.0) && (rp.ps[kd] == 0.1) && (rp.ps[k1] == 0.25) && (rp.ps[k2] == 13.0)
        rp = remake(prob; p = [model.k2 => 14.0])
        @test (rp.ps[kp] == 1.0) && (rp.ps[kd] == 0.1) && (rp.ps[k1] == 0.25) && (rp.ps[k2] == 14.0)
        rp = remake(prob; p = [:k2 => 15.0])
        @test (rp.ps[kp] == 1.0) && (rp.ps[kd] == 0.1) && (rp.ps[k1] == 0.25) && (rp.ps[k2] == 15.0)
    end
end

# Test integrator indexing.
let 
    @test_broken false # NOTE: Multiple problems for `nint`.
    @test_broken false # NOTE: Multiple problems for `jint`.
    for int in deepcopy([oint, sint])
        # Get u values.
        @test int[X] == int[model.X] == int[:X] == 4
        @test int[XY] == int[model.XY] == int[:XY] == 9
        @test int[[XY,Y]] == int[[model.XY,model.Y]] == int[[:XY,:Y]] == [9, 5]
        @test int[(XY,Y)] == int[(model.XY,model.Y)] == int[(:XY,:Y)] == (9, 5)
        @test getu(int, X)(int) == getu(int, model.X)(int) == getu(int, :X)(int) == 4
        @test getu(int, XY)(int) == getu(int, model.XY)(int) == getu(int, :XY)(int) == 9 
        @test getu(int, [XY,Y])(int) == getu(int, [model.XY,model.Y])(int) == getu(int, [:XY,:Y])(int) == [9, 5]  
        @test getu(int, (XY,Y))(int) == getu(int, (model.XY,model.Y))(int) == getu(int, (:XY,:Y))(int) == (9, 5)

        # Set u values.
        int[X] = 20
        @test int[X] == 20
        int[model.X] = 30
        @test int[X] == 30
        int[:X] = 40
        @test int[X] == 40
        setu(int, X)(int, 50)
        @test int[X] == 50
        setu(int, model.X)(int, 60)
        @test int[X] == 60
        setu(int, :X)(int, 70)
        @test int[X] == 70

        # Get p values.
        @test int.ps[kp] == int.ps[model.kp] == int.ps[:kp] == 1.0    
        @test int.ps[[k1,k2]] == int.ps[[model.k1,model.k2]] == int.ps[[:k1,:k2]] == [0.25, 0.5]
        @test int.ps[(k1,k2)] == int.ps[(model.k1,model.k2)] == int.ps[(:k1,:k2)] == (0.25, 0.5)
        @test getp(int, kp)(int) == getp(int, model.kp)(int) == getp(int, :kp)(int) == 1.0
        @test getp(int, [k1,k2])(int) == getp(int, [model.k1,model.k2])(int) == getp(int, [:k1,:k2])(int) == [0.25, 0.5]
        @test getp(int, (k1,k2))(int) == getp(int, (model.k1,model.k2))(int) == getp(int, (:k1,:k2))(int) == (0.25, 0.5)
        
        # Set p values.
        int.ps[kp] = 2.0
        @test int.ps[kp] == 2.0
        int.ps[model.kp] = 3.0
        @test int.ps[kp] == 3.0
        int.ps[:kp] = 4.0
        @test int.ps[kp] == 4.0
        setp(int, kp)(int, 5.0)
        @test int.ps[kp] == 5.0
        setp(int, model.kp)(int, 6.0)
        @test int.ps[kp] == 6.0
        setp(int, :kp)(int, 7.0)
        @test int.ps[kp] == 7.0
    end
end

# Test solve's save_idxs argument.
let 
    for (prob, solver) in zip(deepcopy([oprob, sprob, jprob]), [Tsit5(), ImplicitEM(), SSAStepper()])
        # Save single variable
        @test_broken solve(prob, solver; save_idxs=X)[X][1] == 4
        @test_broken solve(prob, solver; save_idxs=model.X)[X][1] == 4
        @test_broken solve(prob, solver; save_idxs=:X)[X][1] == 4

        # Save observable.
        @test_broken solve(prob, solver; save_idxs=XY)[XY][1] == 9
        @test_broken solve(prob, solver; save_idxs=model.XY)[XY][1] == 9
        @test_broken solve(prob, solver; save_idxs=:XY)[XY][1] == 9

        # Save vector of stuff.
        @test_broken solve(prob, solver; save_idxs=[XY,Y])[[XY,Y]][1] == [9, 5]
        @test_broken solve(prob, solver; save_idxs=[model.XY,model.Y])[[model.XY,model.Y]][1] == [9, 5]
        @test_broken solve(prob, solver; save_idxs=[:XY,:Y])[[:XY,:Y]][1] == [9, 5]
    end
end

nsol[X]

#Tests solution indexing.
let 
    for sol in deepcopy([osol, ssol, jsol])
        # Get u values.
        @test sol[X][1] == sol[model.X][1] == sol[:X][1] == 4
        @test sol[XY][1] == sol[model.XY][1] == sol[:XY][1] == 9
        @test sol[[XY,Y]][1] == sol[[model.XY,model.Y]][1] == sol[[:XY,:Y]][1] == [9, 5]
        @test sol[(XY,Y)][1] == sol[(model.XY,model.Y)][1] == sol[(:XY,:Y)][1] == (9, 5)
        @test getu(sol, X)(sol)[1] == getu(sol, model.X)(sol)[1] == getu(sol, :X)(sol)[1] == 4
        @test getu(sol, XY)(sol)[1] == getu(sol, model.XY)(sol)[1] == getu(sol, :XY)(sol)[1] == 9 
        @test getu(sol, [XY,Y])(sol)[1] == getu(sol, [model.XY,model.Y])(sol)[1] == getu(sol, [:XY,:Y])(sol)[1] == [9, 5]  
        @test getu(sol, (XY,Y))(sol)[1] == getu(sol, (model.XY,model.Y))(sol)[1] == getu(sol, (:XY,:Y))(sol)[1] == (9, 5)       

        # Get u values via idxs and functional call.
        @test osol(0.0; idxs=X) == osol(0.0; idxs=X) == osol(0.0; idxs=X) == 4
        @test osol(0.0; idxs=XY) == osol(0.0; idxs=XY) == osol(0.0; idxs=XY) == 9
        @test_broken osol(0.0; idxs=[model.Y,model.XY]) == osol(0.0; idxs=[model.Y,model.XY]) == osol(0.0; idxs=[model.XY,model.X]) == [9, 5]
        @test_broken osol(0.0; idxs=(:Y,:XY)) == osol(0.0; idxs=(:Y,:XY)) == osol(0.0; idxs=(:XY,:Y)) == (9, 5)

        # Get p values.
        @test sol.ps[kp] == sol.ps[model.kp] == sol.ps[:kp] == 1.0    
        @test sol.ps[[k1,k2]] == sol.ps[[model.k1,model.k2]] == sol.ps[[:k1,:k2]] == [0.25, 0.5]
        @test sol.ps[(k1,k2)] == sol.ps[(model.k1,model.k2)] == sol.ps[(:k1,:k2)] == (0.25, 0.5)
        @test getp(sol, kp)(sol) == getp(sol, model.kp)(sol) == getp(sol, :kp)(sol) == 1.0
        @test getp(sol, [k1,k2])(sol) == getp(sol, [model.k1,model.k2])(sol) == getp(sol, [:k1,:k2])(sol) == [0.25, 0.5]
        @test getp(sol, (k1,k2))(sol) == getp(sol, (model.k1,model.k2))(sol) == getp(sol, (:k1,:k2))(sol) == (0.25, 0.5)
    end
    # handles nonlinear solution differently.
    for sol in deepcopy([nsol])
        # Get u values.
        @test sol[X] == sol[model.X] == sol[:X]
        @test sol[XY] == sol[model.XY][1] == sol[:XY]
        @test sol[[XY,Y]] == sol[[model.XY,model.Y]] == sol[[:XY,:Y]]
        @test_broken sol[(XY,Y)] == sol[(model.XY,model.Y)] == sol[(:XY,:Y)]
        @test getu(sol, X)(sol) == getu(sol, model.X)(sol)[1] == getu(sol, :X)(sol)
        @test getu(sol, XY)(sol) == getu(sol, model.XY)(sol)[1] == getu(sol, :XY)(sol)
        @test getu(sol, [XY,Y])(sol) == getu(sol, [model.XY,model.Y])(sol) == getu(sol, [:XY,:Y])(sol)
        @test_broken getu(sol, (XY,Y))(sol) == getu(sol, (model.XY,model.Y))(sol) == getu(sol, (:XY,:Y))(sol)[1]   

        # Get p values.
        @test sol.ps[kp] == sol.ps[model.kp] == sol.ps[:kp]
        @test sol.ps[[k1,k2]] == sol.ps[[model.k1,model.k2]] == sol.ps[[:k1,:k2]]
        @test sol.ps[(k1,k2)] == sol.ps[(model.k1,model.k2)] == sol.ps[(:k1,:k2)]
        @test getp(sol, kp)(sol) == getp(sol, model.kp)(sol) == getp(sol, :kp)(sol)
        @test getp(sol, [k1,k2])(sol) == getp(sol, [model.k1,model.k2])(sol) == getp(sol, [:k1,:k2])(sol)
        @test getp(sol, (k1,k2))(sol) == getp(sol, (model.k1,model.k2))(sol) == getp(sol, (:k1,:k2))(sol)
    end
end

# Tests plotting.
let 
    @test_broken false # Currently broken for `ssol`.
    for sol in deepcopy([osol, jsol])
        # Single variable.
        @test length(plot(sol; idxs = X).series_list) == 1
        @test length(plot(sol; idxs = XY).series_list) == 1
        @test length(plot(sol; idxs = model.X).series_list) == 1
        @test length(plot(sol; idxs = model.XY).series_list) == 1
        @test length(plot(sol; idxs = :X).series_list) == 1
        @test length(plot(sol; idxs = :XY).series_list) == 1

        # As vector.
        @test length(plot(sol; idxs = [X,Y]).series_list) == 2
        @test length(plot(sol; idxs = [XY,Y]).series_list) == 2
        @test length(plot(sol; idxs = [model.X,model.Y]).series_list) == 2
        @test length(plot(sol; idxs = [model.XY,model.Y]).series_list) == 2
        @test length(plot(sol; idxs = [:X,:Y]).series_list) == 2
        @test length(plot(sol; idxs = [:XY,:Y]).series_list) == 2

        # As tuple.
        @test length(plot(sol; idxs = (X, Y)).series_list) == 1
        @test length(plot(sol; idxs = (XY, Y)).series_list) == 1
        @test length(plot(sol; idxs = (model.X, model.Y)).series_list) == 1
        @test length(plot(sol; idxs = (model.XY, model.Y)).series_list) == 1
        @test length(plot(sol; idxs = (:X, :Y)).series_list) == 1
        @test length(plot(sol; idxs = (:XY, :Y)).series_list) == 1
    end     
end

# Tests solving for various inputs types.
let 
    u0_vals = [X => 1, Y => 0]
    tspan = (0.0, 10.0)
    p_vals = [kp => 1.0, kd => 0.2, k1 => 1.0, k2 => 2.0]

    u0_vals_2 = [model.X => 1, model.Y => 0]
    u0_vals_3 = [:X => 1, :Y => 0]
    p_vals_2 = [model.kp => 1.0, model.kd => 0.2, model.k1 => 1.0, model.k2 => 2.0]
    p_vals_3 = [:kp => 1.0, :kd => 0.2, :k1 => 1.0, :k2 => 2.0]

    oprob_2 = ODEProblem(model, u0_vals_2, tspan, p_vals_2)
    oprob_3 = ODEProblem(model, u0_vals_3, tspan, p_vals_3)
    sprob_2 = SDEProblem(model,u0_vals_2, tspan, p_vals_2)
    sprob_3 = SDEProblem(model,u0_vals_3, tspan, p_vals_3)
    dprob_2 = DiscreteProblem(model, u0_vals_2, tspan, p_vals_2)
    dprob_3 = DiscreteProblem(model, u0_vals_3, tspan, p_vals_3)
    jprob_2 = JumpProblem(model, dprob_2, Direct())
    jprob_3 = JumpProblem(model, dprob_3, Direct())
    nprob_2 = NonlinearProblem(model, u0_vals_2, p_vals_2)
    nprob_3 = NonlinearProblem(model, u0_vals_3, p_vals_3)
    
    @test solve(oprob, Tsit5()) == solve(oprob_2, Tsit5()) == solve(oprob_3, Tsit5())
    @test solve(sprob, ImplicitEM(); seed=1234) == solve(sprob_2, ImplicitEM(); seed=1234) == solve(sprob_3, ImplicitEM(); seed=1234)
    @test solve(jprob, SSAStepper(); seed=1234) == solve(jprob_2, SSAStepper(); seed=1234) == solve(jprob_3, SSAStepper(); seed=1234)
    @test solve(nprob, NewtonRaphson()) == solve(nprob_2, NewtonRaphson()) == solve(nprob_3, NewtonRaphson())
end

[AayushSabharwal]

I believe everything here is solved bar save_idxs, right? Can we rename the issue? It helps to better keep track of things

[ChrisRackauckas]

Yes, save_idxs its own set of issues which needs a longer term fix.

[TorkelE]

Yeah, save_idxs is only superficially related to this I guess, since it is more about implementing something new than filling out some cases. Tons of thanks @AayushSabharwal for fixing all if this!