Syntax for custom cost function for ODE parameter estimation

Hi, I’m trying to write my own cost function for parameter estimation of an ode model.
I am trying to weight the loss by the time t.

Based on the documentation I came up with the following, but I’m unsure about the correct syntax:

function optimization(model,data,constants,dens,options)
# fit function - given model, data, densities, options get best parameters
    @unpack lb,ub,tspan = options
    
    f = (du,u,p,t) -> model(du,u,10 .^ p,t,dens,constants) # optimize in log space
    prob = ODEProblem(f,u₀,tspan,p)
    # cost_function = build_loss_objective(prob,Tsit5(),L2Loss(t,data),
    #                maxiters=10000,verbose=false)
    cost_function = build_loss_objective(prob,Tsit5(),my_loss_function(sol,t,data),
                    maxiters=10000,verbose=false)

    # lb = [-10., -10.] # from 1e-10 to 1e10
    # ub = [5., 5.]
    result = optimize(cost_function,lb,ub,[0.,0.], Fminbox(BFGS()))
end

function my_loss_function(sol,t,data)
   tot_loss = 0.0
   if any((s.retcode != :Success for s in sol))
     tot_loss = Inf
   else
     # calculatisolon for the loss here
     # n = length(t)
     exp = sol(t)
     tot_loss = sum( ((exp .- data) .^ 2) .* t ) # time weighted loss
   end
   tot_loss
end

fit = optimization(model_name,data,constants,dens,c)

Thanks

Use DiffEqFlux for this: DiffEqFlux.jl: Generalized Physics-Informed and Scientific Machine Learning (SciML) · DiffEqFlux.jl. Specifically, this tutorial should be helpful:

https://diffeqflux.sciml.ai/dev/examples/optimization_ode/

Thanks! I’ve rewritten it in the DiffEqFlux style, and it works fine when the dataset is a global variable that the loss function can access, but I am now struggling to pass in a given dataset from my optimization function to the loss function, since the loss has one argument ( p ).

using DifferentialEquations, Flux, Optim, DiffEqFlux, DiffEqSensitivity, Plots

# -------------------------------------------------------------------------
function bimolecular(du,u,p,t,dens)
    # unpack rates and constants
    nᵣ = u[1]
    k₁,k₋₁  = p
    mᵣ,mₗ,A = dens
    # model
    du[1] = dnᵣ = A*k₁*mᵣ*mₗ - k₋₁*nᵣ
end

# --------------------------create dummy data-----------------------------------
f = (du,u,p,t) -> bimolecular(du,u,p,t,dens) # enclose constants

# u0 = [0.0;0.0]
u₀ = [0.0]          # Initial condition
tspan = (0.0,10.0)  # Simulation interval
p = [1e-4,0.5]      # equation parameter. p = [k₁, k₋₁]
dens = [100., 10., 1.]

# Setup the ODE problem, then solve
prob = ODEProblem(f,u₀,tspan,p)
sol = solve(prob)
# sol = solve(prob,Tsit5())
h = plot(sol)

t = collect(range(0,stop=10,length=101))
using RecursiveArrayTools # for VectorOfArray
randomized = VectorOfArray([(sol(t[i]) + .02randn(size(sol,1))) for i in 1:length(t)])
dataset1 = convert(Array,randomized)
scatter!(t,dataset1')


function loss(p) # need to use data that was passed to the optimization function
    tmp_prob = remake(prob,p=p)
    tmp_sol = solve(tmp_prob,saveat=0.1)
    sum(abs2,Array(tmp_sol) - data), tmp_sol
end

function callback(p,l,tmp_sol)
    @show l
    tmp_prob = remake(prob,p=p)
    tmp_sol = solve(tmp_prob,saveat=0.1)
    fig = plot(tmp_sol)
    scatter!(fig,tmp_sol.t,data')
    display(fig)
    false
end



function optimization(model,data,dens,options)
# fit function - given model, data, densities, options get best parameters
    
    lb,ub = options
   
    f = (du,u,p,t) -> model(du,u,p,t,dens) 
    prob = ODEProblem(f,u₀,tspan,p)
    
    res = DiffEqFlux.sciml_train(loss,pinit,ADAM(0.01),cb=callback,maxiters=100)
    res.minimizer
    res = DiffEqFlux.sciml_train(loss,res.minimizer,BFGS(),cb=callback)
end

model_opt2 = (1.0e-6,10.0)
pinit = [2.0e-4,1.0]

# want to use dataset1 for optimization
optimization(bimolecular,dataset1,dens,model_opt2)

It is not. See something like https://diffeqflux.sciml.ai/dev/examples/mnist_neural_ode/ for how the data iterator works.