Zygote.hessian of loss function when defining ODEs with modeling toolkit

General Usage
, , ,
1

I’m doing parameter estimation on an ODE system, which I define using modelingtoolkit. I need to generate hessian of cost function with respect to parameters to access fit. No issue with parameter estimation.

Loss function looks like:

function model_loss(theta)
	_prob = remake(prob,u0=[mn => prescale_means[1],n => prot_val,ms => prescale_means[7], s => prot_val],p=[k1 => theta[1], k2 => theta[2], k3 => theta[3], k4 => theta[4], k5 => theta[5], k6 => theta[6], k7 => theta[7], k8 => theta[8], k9 => theta[9]])
	eko_sol = solve(_prob, Vern7(), saveat=SaveTimes) #forward simulate
	eko_estimated = reduce(hcat, eko_sol.u) #reshape
	mrna_states_eko = vcat(eko_estimated[1,:],eko_estimated[3,:]) #store mRNA states 
	mrna_states_scale = mrna_states_eko/maximum(mrna_states_eko) #scale mRNA states 
	sol = mrna_states_scale
	loss = norm(sqrt.(weight).*( sol -scaled_mrna ),2)

return sol

Then I want to take hessian of cost function with respect to optimal parameters I find.

As,

hessian = Zygote.hessian(z -> model_loss(z), optimized_parameters)

If I don’t use ModelingToolkit to define models this works fine, but I want to use modeling toolkit to programmatically define models.

When I use MTK, I get the error:

ERROR: LoadError: Compiling Tuple{Type{Dict}, Base.Generator{Vector{Equation}, ModelingToolkit.var"#210#211"}}: try/catch is not supported.
Refer to the Zygote documentation for fixes.
https://fluxml.ai/Zygote.jl/latest/limitations

This is a related fix here: Error when differentiating solutions of ModelingToolkit models · Issue #292 · SciML/SciMLBase.jl · GitHub relating to remake

General guidance here for defining problem outside of loss here, but I’m not sure it applies to remake: Frequently Asked Questions · ModelingToolkit.jl

Workaround is shown here: Gradient fails for Dict constructed with a generator or a vector of pairs · Issue #1293 · FluxML/Zygote.jl · GitHub

Parameter estimation works fine, it’s generating the hessian of cost function that causes error.

Goal: i just need to compute hessian of cost function with respect to parameters. I don’t need to use Zygote if something else works.

@ChrisRackauckas

2

Forwarddiff works:

hessian = ForwardDiff.hessian(z -> model_loss(z), optimized_parameters)
1 Like