Problems with accessing ODE parameters after optimisation when modelling with ModelingToolkit

Hello,

I am working with a small non-linear ODE problem to model some protein dynamics data. I am fitting my model to the data and estimating the parameters using Optimization.jl. So far everything worked fine, my problem occurs when I need to retrieve a desired parameter from the optimisation solution.

I built my ODE model with @mtkmodel and defined the ODE problem with ODEProblem. Here is a simpler example of my problem:

@mtkmodel protein_dynamics begin

	@parameters begin
		x1_act
		x1_inact
		x2_act
		x2_inact
	end
	
	@variables begin
		x1(t)
		x2(t)
	end
	
	@equations begin
		D(x1) ~ x1_act  - x1_inact * x1
		D(x2) ~ x2_act  - x2_inact * x2
	end
end

@mtkbuild protein_model = protein_dynamics()

true_param =  [protein_model.x1_act => 0.02, protein_model.x1_inact => 0.2, protein_model.x2_act => 0.04, protein_model.x2_inact => 0.4]
simul_data_prob = ODEProblem(protein_model, [protein_model.x1 => 0.5, protein_model.x2 => 1], (0.0, 200.0), true_param, jac = true)
	
simul_data_sol = solve(simul_data_prob, saveat=1, dtmax=0.5)

# Import the data as array
simul_data_x1 = simul_data_sol[protein_model.x1]
simul_data_x2 = simul_data_sol[protein_model.x2]

When I pass the array of parameters to ODEProblem, it automatically builds an MTKParameters object which shuffles the order of the parameters.

println(simul_data_prob.p)
> ModelingToolkit.MTKParameters{Tuple{Vector{Float64}}, Tuple{}, Tuple{}, Tuple{}, Tuple{}, Nothing, Nothing}(([0.4, 0.02, 0.2, 0.04],), (), (), (), (), nothing, nothing)

This order is kept during the optimisation process and the parameters in the solution array are also in this order.

# Define MSE loss function to estimate parameters
function loss(new_parameters)
	
	# Update the ODE problem with the new set of parameters
	new_ODE_prob = remake(simul_data_prob, p=new_parameters)

	# Solve the new ODE problem 
	new_sol = solve(new_ODE_prob, saveat = 1, dtmax=0.5)
	new_sol_array_x1 = new_sol[protein_model.x1]
	new_sol_array_x2 =  new_sol[protein_model.x2]

	# Determine the loss by computing the sum of squares
	if length(simul_data_x1) == length(new_sol_array_x1)
		loss = sum(abs2, new_sol_array_x1 .- simul_data_x1) + 
		sum(abs2, new_sol_array_x2 .- simul_data_x2)
	else
		loss = 1000
	end

	return loss, new_sol
end

# Define optimization function and differentiation method
optim_f = Optimization.OptimizationFunction((x, p) -> loss(x), Optimization.AutoForwardDiff())
	
# Initial parameter guess
param_guess = [0.1, 0.1, 0.1, 0.1]

# Define the optimisation problem
optim_prob = Optimization.OptimizationProblem(optim_f, param_guess)

# Run the optimisation
results = Optimization.solve(optim_prob, PolyOpt(), maxiters = 100)
println(results.u)
> [0.39999999267413, 0.020000000029211443, 0.19999999991817688, 0.03999999927883234]

I could not figure out how to obtain the index mapping between the initial array and the MTKParameters indexing order although I tried using SymbolicIndexingInterface.jl or SciMLStructures.jl as explained in the FAQ of ModelingToolkit.jl.

Does anyone know how to handle that indexing?

sol.ps[x]

Thanks! Yes, I tried this option but I unfortunately got the following error messages:

println(results.ps[protein_model.x1_act])
ERROR: LoadError: MethodError: no method matching parameter_values(::SciMLBase.NullParameters)
An arithmetic operation was performed on a NullParameters object. This means no parameters were passed
into the AbstractSciMLProblem (e.x.: ODEProblem) but the parameters object `p` was used in an arithmetic
expression. Two common reasons for this issue are:

1. Forgetting to pass parameters into the problem constructor. For example, `ODEProblem(f,u0,tspan)` should
be `ODEProblem(f,u0,tspan,p)` in order to use parameters.

2. Using the wrong function signature. For example, with `ODEProblem`s the function signature is always
`f(du,u,p,t)` for the in-place form or `f(u,p,t)` for the out-of-place form. Note that the `p` argument
will always be in the function signature regardless of if the problem is defined with parameters!



Closest candidates are:
  parameter_values(::Any, ::Any)
   @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:16
  parameter_values(::SciMLBase.AbstractJumpProblem)
   @ SciMLBase ~/.julia/packages/SciMLBase/Dwomw/src/problems/problem_interface.jl:12
  parameter_values(::ModelingToolkit.MTKParameters, ::ModelingToolkit.ParameterIndex)
   @ ModelingToolkit ~/.julia/packages/ModelingToolkit/Mxj1Q/src/systems/parameter_buffer.jl:216
  ...

Stacktrace:
 [1] parameter_values(prob::SciMLBase.NullParameters, i::Nothing)
   @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:16
 [2] parameter_values(A::SciMLBase.OptimizationSolution{…}, i::Nothing)
   @ SciMLBase ~/.julia/packages/SciMLBase/Dwomw/src/solutions/solution_interface.jl:44
 [3] _getter
   @ ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:117 [inlined]
 [4] getter
   @ ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:139 [inlined]
 [5] getindex(::SymbolicIndexingInterface.ParameterIndexingProxy{SciMLBase.OptimizationSolution{…}}, ::Num)
   @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing_proxy.jl:14
 [6] top-level scope
   @ ~/Documents/Master/MasterThesis/RhoAModelling/Notebooks/parameters.jl:84
 [7] include(fname::String)
   @ Base.MainInclude ./client.jl:489
 [8] top-level scope
   @ REPL[2]:1

Or on the ODE solution:

sol = solve(remake(simul_data_prob, p=results.u), saveat=0.5, dtmax = 0.5)
println(sol.ps[protein_model.x1_act])
ERROR: LoadError: ArgumentError: invalid index: ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, Tuple{Int64, Int64}}(SciMLStructures.Tunable(), (1, 2)) of type ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, Tuple{Int64, Int64}}
Stacktrace:
  [1] to_index(i::ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, Tuple{Int64, Int64}})
    @ Base ./indices.jl:300
  [2] to_index(A::Vector{Float64}, i::ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, Tuple{Int64, Int64}})
    @ Base ./indices.jl:277
  [3] _to_indices1(A::Vector{Float64}, inds::Tuple{Base.OneTo{…}}, I1::ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, Tuple{…}})
    @ Base ./indices.jl:359
  [4] to_indices
    @ ./indices.jl:354 [inlined]
  [5] to_indices
    @ ./indices.jl:345 [inlined]
  [6] getindex
    @ ./abstractarray.jl:1291 [inlined]
  [7] parameter_values(arr::Vector{Float64}, i::ModelingToolkit.ParameterIndex{SciMLStructures.Tunable, Tuple{Int64, Int64}})
    @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:15
  [8] parameter_values(A::ODESolution{…}, i::ModelingToolkit.ParameterIndex{…})
    @ SciMLBase ~/.julia/packages/SciMLBase/Dwomw/src/solutions/solution_interface.jl:44
  [9] (::SymbolicIndexingInterface.var"#_getter#9"{…})(::SymbolicIndexingInterface.Timeseries, prob::ODESolution{…})
    @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:120
 [10] (::SymbolicIndexingInterface.var"#getter#10"{…})(::ODESolution{…})
    @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing.jl:139
 [11] getindex(::SymbolicIndexingInterface.ParameterIndexingProxy{ODESolution{…}}, ::Num)
    @ SymbolicIndexingInterface ~/.julia/packages/SymbolicIndexingInterface/Pz8Cf/src/parameter_indexing_proxy.jl:14
 [12] top-level scope
    @ ~/Documents/Master/MasterThesis/RhoAModelling/Notebooks/parameters.jl:88
 [13] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [14] top-level scope
    @ REPL[4]:1

Your optimization problem doesn’t have parameters, the values you looking for are part of the state. So println(results[protein_model.x1_act]).

The second one, we need to document how to do this better.

@cryptic.ax a good example

I wasn’t sure you were making reference to the optimisation solution or ODE solution, so I tried both. Unfortunately, the piece of code is still not working for me:

println(results[protein_model.x1_act])
ERROR: LoadError: Indexing symbol x1_act is unknown.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] DEFAULT_OBSERVED_NO_TIME(sym::Num, u::Vector{Float64}, p::SciMLBase.NullParameters)
   @ SciMLBase ~/.julia/packages/SciMLBase/Dwomw/src/scimlfunctions.jl:177
 [3] (::SciMLBase.var"#187#188"{OptimizationFunction{…}, Num})(::Vector{Float64}, ::Vararg{Any})
   @ SciMLBase ~/.julia/packages/SciMLBase/Dwomw/src/scimlfunctions.jl:4072
 [4] getindex(A::SciMLBase.OptimizationSolution{…}, sym::Num)
   @ SciMLBase ~/.julia/packages/SciMLBase/Dwomw/src/solutions/solution_interface.jl:72
 [5] top-level scope
   @ ~/Documents/Master/MasterThesis/RhoAModelling/Notebooks/parameters.jl:89
 [6] include(fname::String)
   @ Base.MainInclude ./client.jl:489
 [7] top-level scope
   @ REPL[4]:1

Seems like observed needs a bit of work on the optimization. We’ll take a look over the coming week.