15. Parameter Estimation with Noisy Data

Introduction

In this example, we demonstrate the filtering power of the Prediction Error Method (PEM). We generate data from a true Hodgkin-Huxley neuron, add realistic Gaussian noise to the voltage trace, and then attempt to recover the conductances. The PEM observer channel acts as a feedback controller, absorbing the noise while allowing the optimizer to find the true underlying parameters without overfitting the noise.

using MTKNeuralToolkit
using SymbolicIndexingInterface: getu
using MTKNeuralToolkit.HodgkinHuxley: SodiumChannel, PotassiumChannel, LeakChannel
using ModelingToolkit: mtkcompile, @named
using OrdinaryDiffEq
using OrdinaryDiffEqRosenbrock
using Optimization
using OptimizationOptimJL
using SciMLStructures: Tunable, canonicalize, replace
using SymbolicIndexingInterface: parameter_values, setp
using PreallocationTools
using DataInterpolations
using SciMLBase
using Plots
using Markdown
using Random

1. Build the True System & Generate Noisy Data

top = Scalar()

function build_hh_neuron(name::Symbol; gNa=120.0, gK=36.0, gleak=0.3, pem=false, itps=nothing, K=1.0)
    @named cap  = Capacitor(topology=top, C=1.0)
    @named na   = SodiumChannel(topology=top, g=gNa)
    @named k    = PotassiumChannel(topology=top, g=gK)
    @named leak = LeakChannel(topology=top, g=gleak)
    channels = [na, k, leak]
    if pem
        @named pem_ch = PEMObservationChannel(itps=itps, K_init=K, topology=top)
        push!(channels, pem_ch)
    end
    return build_compartment(cap, channels; name=name, V_init=-65.0, topology=top)
end

true_gNa = 120.0
true_gK  = 36.0
true_gleak = 0.3

true_neuron = build_hh_neuron(:true_neuron; gNa=true_gNa, gK=true_gK, gleak=true_gleak)

drivers = [(1, 8.0)]
true_net = build_acausal_network([true_neuron]; drivers=drivers, name=:true_net)
END_TIME= 100.0


true_sys = mtkcompile(true_net.sys)
true_prob = ODEProblem(true_sys, [], (0.0, END_TIME))
timesteps = 0.0:0.1:END_TIME

true_sol = solve(true_prob, Rodas5(); saveat=timesteps)
V_data_clean = true_sol[true_sys.true_neuron.cap.v]

1001-element Vector{Float64}:
 -65.0
 -64.22591619267946
 -63.49089980908908
 -62.78551666686095
 -62.10164954963334
 -61.43161497634572
 -60.767610613029596
 -60.10127329104882
 -59.423257082517694
 -58.722544449144145
   ⋮
  10.031725507089753
   4.687598428791057
  -0.5743779076458245
  -5.700204926125791
 -10.665492866763016
 -15.4686757127778
 -20.12347941188918
 -24.65847753171579
 -29.127817545972867

Add 3 mV Gaussian noise to the data

Random.seed!(42)
noise_level = 3.0
V_data_noisy = V_data_clean .+ noise_level .* randn(length(V_data_clean))
itp_V = LinearInterpolation(V_data_noisy, timesteps)

LinearInterpolation with 1001 points
┌───────┬──────────┐
│     t │        u │
├───────┼──────────┤
│   0.0 │ -66.0901 │
│   0.1 │ -63.4707 │
│   0.2 │ -64.4359 │
│   0.3 │ -63.7193 │
│   0.4 │ -59.6527 │
│   0.5 │ -60.0014 │
│   0.6 │ -63.3463 │
│   0.7 │ -64.5091 │
│     ⋮ │        ⋮ │
│  99.3 │  3.27497 │
│  99.4 │ -1.30911 │
│  99.5 │ -4.74927 │
│  99.6 │ -8.18907 │
│  99.7 │ -21.2467 │
│  99.8 │ -17.1929 │
│  99.9 │ -25.3975 │
│ 100.0 │ -29.7477 │
└───────┴──────────┘
    985 rows omitted

2. Setup the PEM Optimization Problem

We create a model neuron with terrible initial guesses and attach a PEM observer to the noisy data.

guess_gNa = 10.0
guess_gK  = 100.0
guess_gleak = 10.3

fit_neuron = build_hh_neuron(:fit_neuron; gNa=guess_gNa, gK=guess_gK, gleak=guess_gleak, pem=true, itps=[itp_V], K=2.0)
fit_net = build_acausal_network([fit_neuron]; drivers=drivers, name=:fit_net)
fit_sys = mtkcompile(fit_net.sys)
fit_prob = ODEProblem(fit_sys, [], (0.0, END_TIME))

gNa_sym = fit_sys.fit_neuron.na.g
gK_sym  = fit_sys.fit_neuron.k.g
gleak_sym = fit_sys.fit_neuron.leak.g

setter = setp(fit_prob, [gNa_sym, gK_sym, gleak_sym])
diffcache = DiffCache(copy(canonicalize(Tunable(), parameter_values(fit_prob))[1]))

v_getter = getu(fit_prob, fit_sys.fit_neuron.cap.v)
i_pem_getter = getu(fit_prob, fit_sys.fit_neuron.pem_ch.i)

SymbolicIndexingInterface.TimeDependentObservedFunction{SymbolicIndexingInterface.ContinuousTimeseries, ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :__argₛᵧₘ1249670312406203976), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xe2969c25, 0xf1bea75d, 0xd81461b5, 0x899298c5, 0xc90633a4), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x1, :x2, :x3, :x4), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xb896e553, 0xa118fcf5, 0xbfeb9719, 0xa096e58a, 0x357542a4), Nothing}}, true}(SymbolicIndexingInterface.ContinuousTimeseries(), ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :__argₛᵧₘ1249670312406203976), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xe2969c25, 0xf1bea75d, 0xd81461b5, 0x899298c5, 0xc90633a4), Nothing}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x1, :x2, :x3, :x4), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xb896e553, 0xa118fcf5, 0xbfeb9719, 0xa096e58a, 0x357542a4), Nothing}}(RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__mtk_arg_1, :___mtkparameters___, :__argₛᵧₘ1249670312406203976), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xe2969c25, 0xf1bea75d, 0xd81461b5, 0x899298c5, 0xc90633a4), Nothing}(nothing), RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:x1, :x2, :x3, :x4), ModelingToolkitBase.var"#_RGF_ModTag", ModelingToolkitBase.var"#_RGF_ModTag", (0xb896e553, 0xa118fcf5, 0xbfeb9719, 0xa096e58a, 0x357542a4), Nothing}(nothing)))

3. Define Loss Function & Optimize

We calculate the loss based on the PEM observer current. Instead of minimizing the tracking error (which the controller achieves regardless of parameters), we minimize the effort of the controller. This forces the optimizer to find the underlying parameters that naturally generate the data.

3. Define Loss Function & Optimize

We calculate a multi-objective loss based on both the tracking error (voltage) and the observer effort (current). This forces the system to track the data while penalizing the controller from "cheating" to force bad parameters to fit.

function loss(x, p)
    prob, timesteps, V_data_noisy, setter, diffcache, v_getter, i_pem_getter = p

    ps = parameter_values(prob)
    buffer = get_tmp(diffcache, x)
    copyto!(buffer, canonicalize(Tunable(), ps)[1])
    ps = replace(Tunable(), ps, buffer)
    setter(ps, x)
    newprob = remake(prob; p=ps)
    sol = solve(newprob, Rodas5(); saveat=timesteps, reltol=1e-8, abstol=1e-8)

    if !SciMLBase.successful_retcode(sol.retcode)
        return Inf
    end

    V_fit = v_getter(sol)
    I_pem = i_pem_getter(sol)

    #Multi-objective cost
    tracking_error = sum(abs2, V_fit .- V_data_noisy) / length(V_data_noisy)
    observer_effort = sum(abs2, I_pem) / length(I_pem)

    alpha = 1.0
    beta = 1.0

    return alpha * tracking_error + beta * observer_effort
end

loss (generic function with 1 method)

Tuple must also have exactly 7 items:

opt_params = (fit_prob, timesteps, V_data_noisy, setter, diffcache, v_getter, i_pem_getter)
adtype = AutoForwardDiff()
optfn = OptimizationFunction(loss, adtype)
optprob = OptimizationProblem(optfn, [guess_gNa, guess_gK, guess_gleak], opt_params)

OptimizationProblem. In-place: true
u0: 3-element Vector{Float64}:
  10.0
 100.0
  10.3

4. Optimize and Plot

To avoid recompiling new systems for the free-running simulations, we simply reuse the compiled fit_sys and set the PEM controller gain (K) to 0.0 to disable the observer.

println("Starting optimization...")

Starting optimization...

(Don't forget to add lb and ub to optprob in Section 3 if you haven't already!)

res = solve(optprob, BFGS(); maxiters=1000)

retcode: Success
u: 3-element Vector{Float64}:
 113.16727777251688
  33.004973853658065
   0.4402750975226786

Assuming the PEM channel has a parameter named K. (If the symbol name is different in your package, just adjust this line)

K_sym = fit_sys.fit_neuron.pem_ch.K
K_setter = setp(fit_prob, [K_sym])

SymbolicIndexingInterface.ParameterHookWrapper{SymbolicIndexingInterface.MultipleSetters{Vector{SymbolicIndexingInterface.SetParameterIndex{ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, Int64}}}}, Vector{Symbolics.Num}}(SymbolicIndexingInterface.MultipleSetters{Vector{SymbolicIndexingInterface.SetParameterIndex{ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, Int64}}}}(SymbolicIndexingInterface.SetParameterIndex{ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, Int64}}[SymbolicIndexingInterface.SetParameterIndex{ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, Int64}}(ModelingToolkitBase.ParameterIndex{SciMLStructures.Tunable, Int64}(SciMLStructures.Tunable(), 8, false))]), Symbolics.Num[fit_neuron₊pem_ch₊K])

–- 1. Simulate with recovered parameters (WITH PEM) to get observer current –-

opt_ps = parameter_values(fit_prob)
opt_buffer = copy(canonicalize(Tunable(), opt_ps)[1])
opt_ps = replace(Tunable(), opt_ps, opt_buffer)
setter(opt_ps, res.u) # set recovered parameters
opt_prob_pem = remake(fit_prob; p=opt_ps)
opt_sol_pem = solve(opt_prob_pem, Rodas5(); saveat=timesteps, reltol=1e-6, abstol=1e-6)

retcode: Success
Interpolation: 1st order linear
t: 1001-element Vector{Float64}:
   0.0
   0.1
   0.2
   0.3
   0.4
   0.5
   0.6
   0.7
   0.8
   0.9
   ⋮
  99.2
  99.3
  99.4
  99.5
  99.6
  99.7
  99.8
  99.9
 100.0
u: 1001-element Vector{Vector{Float64}}:
 [0.052, 0.596, 0.317, -65.0]
 [0.05331216439387372, 0.5958239223128492, 0.31713311174661074, -64.09691276427284]
 [0.05621773635178602, 0.5952573720112642, 0.317529033157161, -63.26733232815712]
 [0.05975305175826559, 0.5943921259309658, 0.31812192082543556, -62.63955603320987]
 [0.06390906682033305, 0.5932079572503919, 0.31892154181157795, -61.72081610499736]
 [0.06944918826250716, 0.5915291086559609, 0.3200373591177592, -60.670938209501436]
 [0.07543993469496087, 0.5894628513997942, 0.32139274361881565, -60.15964379408399]
 [0.08026745094376828, 0.5873100710534978, 0.32279342612632383, -60.12367748952423]
 [0.08356008562513449, 0.5852116090952534, 0.3241504102800607, -60.26107927744527]
 [0.08607611949532325, 0.5831185440460898, 0.3254949967216524, -59.794466994119354]
 ⋮
 [0.9875021260635225, 0.14343376739595878, 0.6854610088512542, 10.98743996665061]
 [0.9868767406529534, 0.1301036322541978, 0.7014170645683921, 5.576905719870741]
 [0.9840338945602634, 0.11818186101721673, 0.7145889387932626, 0.3181570630878811]
 [0.9790441624286019, 0.10758586460127327, 0.7253338119894982, -4.709968218373437]
 [0.9716335228761126, 0.09825349968010146, 0.7339495195802295, -9.551956054548583]
 [0.9609713201943931, 0.09016515702875699, 0.7406119125581689, -14.75970239476487]
 [0.9459687841488083, 0.08332851827517906, 0.7454654547299101, -19.23142779358623]
 [0.9265337711834237, 0.07766401850093066, 0.7488463828538096, -23.692251307088352]
 [0.900509908539332, 0.07322880289578418, 0.7507704928158713, -28.23116268433953]

–- 2. Pure simulation for INITIAL GUESS (No PEM) –-

init_ps = parameter_values(fit_prob)
init_buffer = copy(canonicalize(Tunable(), init_ps)[1])
init_ps = replace(Tunable(), init_ps, init_buffer)
setter(init_ps, [guess_gNa, guess_gK, guess_gleak])
K_setter(init_ps, [0.0]) # disable PEM controller
init_prob_free = remake(fit_prob; p=init_ps)
init_eval_sol = solve(init_prob_free, Rodas5(); saveat=timesteps)

retcode: Success
Interpolation: 1st order linear
t: 1001-element Vector{Float64}:
   0.0
   0.1
   0.2
   0.3
   0.4
   0.5
   0.6
   0.7
   0.8
   0.9
   ⋮
  99.2
  99.3
  99.4
  99.5
  99.6
  99.7
  99.8
  99.9
 100.0
u: 1001-element Vector{Vector{Float64}}:
 [0.052, 0.596, 0.317, -65.0]
 [0.061290858530422494, 0.5942952843383317, 0.318137305411624, -58.70817014124696]
 [0.07730047544215801, 0.590609742063318, 0.3204915318624678, -56.702232466545524]
 [0.09253428914116847, 0.5862459234196883, 0.32322088468856697, -56.08879527296019]
 [0.10502266115176447, 0.5817235920244812, 0.3260211377312959, -55.926509323552615]
 [0.1146688286712826, 0.577220825406562, 0.32879184097563885, -55.910557539390936]
 [0.12189796058700975, 0.5727952528923334, 0.3315011249766164, -55.94299551617634]
 [0.127203089557158, 0.5684649201604384, 0.3341393664088835, -55.99270195382879]
 [0.1310163984823134, 0.5642352991472477, 0.3367039957035816, -56.04956671005895]
 [0.13368349618210665, 0.5601076560512765, 0.3391947897373976, -56.11013491529922]
 ⋮
 [0.10648816061853868, 0.3779511514913714, 0.4151273978453165, -58.8175980514145]
 [0.10648816061768569, 0.37795114934422624, 0.41512739784167113, -58.817598051485795]
 [0.10648816061684603, 0.3779511472307147, 0.41512739783808306, -58.81759805155598]
 [0.10648816061601968, 0.37795114515082723, 0.4151273978345522, -58.81759805162506]
 [0.10648816061520663, 0.37795114310457045, 0.41512739783107855, -58.817598051693025]
 [0.1064881606144069, 0.37795114109196726, 0.4151273978276623, -58.81759805175989]
 [0.10648816061362049, 0.3779511391130565, 0.41512739782430336, -58.81759805182563]
 [0.10648816061284742, 0.37795113716789325, 0.4151273978210019, -58.817598051890265]
 [0.10648816061208773, 0.3779511352565488, 0.41512739781775804, -58.81759805195378]

–- 3. Pure simulation for RECOVERED PARAMETERS (No PEM) –-

opt_gNa, opt_gK, opt_gleak = res.u
fit_ps = parameter_values(fit_prob)
fit_buffer = copy(canonicalize(Tunable(), fit_ps)[1])
fit_ps = replace(Tunable(), fit_ps, fit_buffer)
setter(fit_ps, [opt_gNa, opt_gK, opt_gleak])
K_setter(fit_ps, [0.0]) # disable PEM controller
fit_prob_free = remake(fit_prob; p=fit_ps)
fit_eval_sol = solve(fit_prob_free, Rodas5(); saveat=timesteps)

retcode: Success
Interpolation: 1st order linear
t: 1001-element Vector{Float64}:
   0.0
   0.1
   0.2
   0.3
   0.4
   0.5
   0.6
   0.7
   0.8
   0.9
   ⋮
  99.2
  99.3
  99.4
  99.5
  99.6
  99.7
  99.8
  99.9
 100.0
u: 1001-element Vector{Vector{Float64}}:
 [0.052, 0.596, 0.317, -65.0]
 [0.05341867593626602, 0.5958039056521621, 0.3171467000773709, -64.05826525880595]
 [0.05637176627527607, 0.5952224864124358, 0.31755251363382514, -63.17469716500626]
 [0.06034231347219115, 0.5942669066420377, 0.3182052197606379, -62.33656943792283]
 [0.0650341602586268, 0.592944150719722, 0.3190951471448403, -61.53255719107533]
 [0.07028243573542317, 0.5912569352315388, 0.32021526614609785, -60.75165542944803]
 [0.07600659634955867, 0.5892030770041314, 0.3215615484615198, -59.98250219783454]
 [0.08218597586455446, 0.5867746441424206, 0.3231333603904243, -59.21280093996116]
 [0.08884867341491619, 0.5839561849916575, 0.32493438767879923, -58.4286281737858]
 [0.09606962787671913, 0.5807228857377956, 0.3269733100056945, -57.61335787227967]
 ⋮
 [0.2076972876536706, 0.37563265682534774, 0.43137477407149527, -49.451062321387]
 [0.22186440164292248, 0.3700823989830402, 0.4345678467650881, -48.05590609718906]
 [0.24026190716620024, 0.3635086317719034, 0.43833444085404744, -46.189549070721704]
 [0.26494353100073925, 0.35545299466374725, 0.44291467885378266, -43.59242977739668]
 [0.299417495426062, 0.3451144549446373, 0.44871797878718717, -39.81261445978859]
 [0.34978421310523233, 0.33117382874171514, 0.4564820183115967, -34.03314721750796]
 [0.42610150234788735, 0.3118833644559568, 0.46754539580418064, -24.762201689650478]
 [0.5401587947081856, 0.2872554440620006, 0.4841766387945265, -10.148045950422103]
 [0.6863255824272171, 0.26106443436671883, 0.5088173196720875, 8.729110523337598]

(Plotting code remains exactly the same)

p1 = plot(timesteps, V_data_noisy, label="Noisy Target", color=:gray, lw=1, alpha=0.8)
plot!(p1, timesteps, V_data_clean, label="True Clean", color=:black, lw=2)
plot!(p1, timesteps, init_eval_sol[fit_sys.fit_neuron.cap.v], label="Initial Guess (Free-Running)", ls=:dot, lw=2, color=:blue)
plot!(p1, timesteps, fit_eval_sol[fit_sys.fit_neuron.cap.v], label="Fit (Free-Running)", ls=:dash, lw=2, color=:red)
title!("Voltage Trace Recovery (Noisy Data)")

Plot the observer current.

I_obs = i_pem_getter(opt_sol_pem)
p2 = plot(timesteps, I_obs, label="Observer Current", color=:red, lw=1.5)
title!("PEM Observer Current (Absorbing Noise)")
hline!([0.0], color=:gray, ls=:dot, label="0")

p = plot(p1, p2, layout=(2,1), size=(800, 700), legend=:outertop)
xlabel!(p, "Time (ms)")
ylabel!(p, "V (mV) / I (nA)")

5. Parameter Comparison

Markdown.parse("""
| Parameter | True Value | Initial Guess | Recovered Value |
|-----------|------------|---------------|-----------------|
| gNa       | $true_gNa  | $guess_gNa    | $(round(opt_gNa, digits=3)) |
| gK        | $true_gK   | $guess_gK     | $(round(opt_gK, digits=3)) |
| gleak     | $true_gleak| $guess_gleak  | $(round(opt_gleak, digits=3)) |
""")

This page was generated using Literate.jl.