# Parameter estimation of FitzHugh-Nagumo model using optimisation methods

using ParameterizedFunctions, OrdinaryDiffEq, DiffEqParamEstim
gr(fmt=:png)

Plots.GRBackend()

loc_bounds = Tuple{Float64,Float64}[(0, 1), (0, 1), (0, 1), (0, 1)]
glo_bounds = Tuple{Float64,Float64}[(0, 5), (0, 5), (0, 5), (0, 5)]
loc_init = [0.5,0.5,0.5,0.5]
glo_init = [2.5,2.5,2.5,2.5]

fitz = @ode_def FitzhughNagumo begin
dv = v - v^3/3 -w + l
dw = τinv*(v +  a - b*w)
end a b τinv l

p = [0.7,0.8,0.08,0.5]              # Parameters used to construct the dataset
r0 = [1.0; 1.0]                     # initial value
tspan = (0.0, 30.0)                 # sample of 3000 observations over the (0,30) timespan
prob = ODEProblem(fitz, r0, tspan,p)
tspan2 = (0.0, 3.0)                 # sample of 300 observations with a timestep of 0.01
prob_short = ODEProblem(fitz, r0, tspan2,p)

dt = 30.0/3000
tf = 30.0
tinterval = 0:dt:tf
t  = collect(tinterval)

h = 0.01
M = 300
tstart = 0.0
tstop = tstart + M * h
tinterval_short = 0:h:tstop
t_short = collect(tinterval_short)

#Generate Data
data_sol_short = solve(prob_short,Vern9(),saveat=t_short,reltol=1e-9,abstol=1e-9)
data_short = convert(Array, data_sol_short) # This operation produces column major dataset obs as columns, equations as rows
data_sol = solve(prob,Vern9(),saveat=t,reltol=1e-9,abstol=1e-9)
data = convert(Array, data_sol)


#### Plot of the solution

##### Short Solution
plot(data_sol_short)

##### Longer Solution
plot(data_sol)


## Local Solution from the short data set

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short)
res1 = bboptimize(obj_short;SearchRange = glo_bounds, MaxSteps = 7e3)

Starting optimization with optimizer BlackBoxOptim.DiffEvoOpt{BlackBoxOptim
daptiveDiffEvoRandBin{3},BlackBoxOptim.RandomBound{BlackBoxOptim.RangePerDi
mSearchSpace}}
0.00 secs, 0 evals, 0 steps
0.50 secs, 1736 evals, 1648 steps, improv/step: 0.263 (last = 0.2633), fitn
ess=0.046172657
1.00 secs, 3648 evals, 3560 steps, improv/step: 0.201 (last = 0.1464), fitn
ess=0.021984202
1.50 secs, 5714 evals, 5626 steps, improv/step: 0.170 (last = 0.1171), fitn
ess=0.000380267

Optimization stopped after 7001 steps and 1.836867094039917 seconds
Termination reason: Max number of steps (7000) reached
Steps per second = 3811.3808139500925
Function evals per second = 3859.288471660078
Improvements/step = 0.16142857142857142
Total function evaluations = 7089

Best candidate found: [0.335025, 0.749594, 0.119288, 0.499943]

Fitness: 0.000057967

# Lower tolerance could lead to smaller fitness (more accuracy)

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short,reltol=1e-9)
res1 = bboptimize(obj_short;SearchRange = glo_bounds, MaxSteps = 7e3)

Starting optimization with optimizer BlackBoxOptim.DiffEvoOpt{BlackBoxOptim
daptiveDiffEvoRandBin{3},BlackBoxOptim.RandomBound{BlackBoxOptim.RangePerDi
mSearchSpace}}
0.00 secs, 0 evals, 0 steps
0.50 secs, 2016 evals, 1906 steps, improv/step: 0.220 (last = 0.2198), fitn
ess=0.096210176
1.00 secs, 4036 evals, 3927 steps, improv/step: 0.168 (last = 0.1197), fitn
ess=0.024551856
1.50 secs, 6071 evals, 5963 steps, improv/step: 0.149 (last = 0.1130), fitn
ess=0.000148160

Optimization stopped after 7001 steps and 1.752614974975586 seconds
Termination reason: Max number of steps (7000) reached
Steps per second = 3994.6024083798125
Function evals per second = 4055.0834618419262
Improvements/step = 0.14842857142857144
Total function evaluations = 7107

Best candidate found: [0.791543, 0.883046, 0.0798279, 0.500219]

Fitness: 0.000028308

# Change in tolerance makes it worse

obj_short = build_loss_objective(prob_short,Vern9(),L2Loss(t_short,data_short),tstops=t_short,reltol=1e-9,abstol=1e-9)
res1 = bboptimize(obj_short;SearchRange = glo_bounds, MaxSteps = 7e3)

Starting optimization with optimizer BlackBoxOptim.DiffEvoOpt{BlackBoxOptim
daptiveDiffEvoRandBin{3},BlackBoxOptim.RandomBound{BlackBoxOptim.RangePerDi
mSearchSpace}}
0.00 secs, 0 evals, 0 steps
0.50 secs, 1649 evals, 1529 steps, improv/step: 0.262 (last = 0.2616), fitn
ess=0.096497628
1.00 secs, 3329 evals, 3210 steps, improv/step: 0.190 (last = 0.1249), fitn
ess=0.003551431
1.50 secs, 4964 evals, 4845 steps, improv/step: 0.163 (last = 0.1095), fitn
ess=0.000142969
2.00 secs, 6545 evals, 6426 steps, improv/step: 0.154 (last = 0.1278), fitn
ess=0.000001087

Optimization stopped after 7001 steps and 2.184602975845337 seconds
Termination reason: Max number of steps (7000) reached
Steps per second = 3204.7013015218236
Function evals per second = 3259.1734419133527
Improvements/step = 0.15528571428571428
Total function evaluations = 7120

Best candidate found: [0.667771, 0.790215, 0.0818976, 0.500033]

Fitness: 0.000001087

# using the moe accurate Vern9() reduces the fitness marginally and leads to some increase in time taken


## Using NLopt

#### Global Optimisation

obj_short = build_loss_objective(prob_short,Vern9(),L2Loss(t_short,data_short),tstops=t_short,reltol=1e-9,abstol=1e-9)

opt = Opt(:GN_ORIG_DIRECT_L, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

2.905493 seconds (12.94 M allocations: 1.156 GiB, 10.02% gc time)
(0.11016600768053908, [0.192044, 1.13169, 1.11111, 0.509578], :XTOL_REACHED
)

opt = Opt(:GN_CRS2_LM, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

2.986420 seconds (13.41 M allocations: 1.199 GiB, 10.24% gc time)
(3.172386310240739e-14, [0.700002, 0.8, 0.0799999, 0.5], :MAXEVAL_REACHED)

opt = Opt(:GN_ISRES, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

2.963639 seconds (13.41 M allocations: 1.199 GiB, 9.97% gc time)
(0.014540158883062867, [4.88552, 3.82717, 0.0415995, 0.508172], :MAXEVAL_RE
ACHED)

opt = Opt(:GN_ESCH, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

2.970191 seconds (13.41 M allocations: 1.199 GiB, 10.04% gc time)
(0.007721132655303023, [3.07127, 2.78539, 0.065985, 0.504515], :MAXEVAL_REA
CHED)


Now local optimization algorithms are used to check the global ones, these use the local constraints, different intial values and time step

opt = Opt(:LN_BOBYQA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

0.720314 seconds (3.18 M allocations: 290.917 MiB, 9.80% gc time)
(5.28611876473971e-24, [0.7, 0.8, 0.08, 0.5], :SUCCESS)

opt = Opt(:LN_NELDERMEAD, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

0.296068 seconds (1.24 M allocations: 113.667 MiB, 11.57% gc time)
(8.965505337540987e-5, [1.0, 1.0, 0.0735509, 0.500405], :XTOL_REACHED)

opt = Opt(:LD_SLSQP, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

0.584728 seconds (2.44 M allocations: 216.328 MiB, 10.25% gc time)
(3.732246654524026e-14, [0.699986, 0.799998, 0.080001, 0.5], :XTOL_REACHED)

opt = Opt(:LN_COBYLA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

2.998466 seconds (13.41 M allocations: 1.199 GiB, 10.08% gc time)
(0.0007533012926827143, [0.182327, 0.834403, 0.195088, 0.500362], :MAXEVAL_
REACHED)

opt = Opt(:LN_NEWUOA_BOUND, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

0.125620 seconds (311.13 k allocations: 28.479 MiB, 4.46% gc time)
(0.0003911847251197023, [0.320842, 0.439003, 0.078888, 0.499177], :SUCCESS)

opt = Opt(:LN_PRAXIS, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

0.300467 seconds (1.25 M allocations: 114.771 MiB, 9.31% gc time)
(4.890413512906367e-25, [0.7, 0.8, 0.08, 0.5], :SUCCESS)

opt = Opt(:LN_SBPLX, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

3.008117 seconds (13.41 M allocations: 1.199 GiB, 10.07% gc time)
(8.350546444056411e-15, [0.700007, 0.800002, 0.0799996, 0.5], :MAXEVAL_REAC
HED)

opt = Opt(:LD_MMA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj_short.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

26.731639 seconds (120.39 M allocations: 10.780 GiB, 10.16% gc time)
(9.930463767834086e-5, [0.231856, 0.706134, 0.130932, 0.49972], :MAXEVAL_RE
ACHED)


### Now the longer problem is solved for a global solution

Vern9 solver with reltol=1e-9 and abstol=1e-9 is used and the dataset is increased to 3000 observations per variable with the same integration time step of 0.01.

obj = build_loss_objective(prob,Vern9(),L2Loss(t,data),tstops=t,reltol=1e-9,abstol=1e-9)
res1 = bboptimize(obj;SearchRange = glo_bounds, MaxSteps = 4e3)

Starting optimization with optimizer BlackBoxOptim.DiffEvoOpt{BlackBoxOptim
daptiveDiffEvoRandBin{3},BlackBoxOptim.RandomBound{BlackBoxOptim.RangePerDi
mSearchSpace}}
0.00 secs, 0 evals, 0 steps
0.51 secs, 158 evals, 94 steps, improv/step: 0.404 (last = 0.4043), fitness
=2375.077980855
1.01 secs, 325 evals, 220 steps, improv/step: 0.395 (last = 0.3889), fitnes
s=1629.212930592
1.51 secs, 487 evals, 376 steps, improv/step: 0.356 (last = 0.3013), fitnes
s=940.066811899
2.02 secs, 651 evals, 540 steps, improv/step: 0.320 (last = 0.2378), fitnes
s=940.066811899
2.52 secs, 817 evals, 706 steps, improv/step: 0.302 (last = 0.2410), fitnes
s=940.066811899
3.02 secs, 985 evals, 874 steps, improv/step: 0.279 (last = 0.1845), fitnes
s=675.366681030
3.52 secs, 1153 evals, 1043 steps, improv/step: 0.251 (last = 0.1065), fitn
ess=675.366681030
4.02 secs, 1323 evals, 1213 steps, improv/step: 0.232 (last = 0.1118), fitn
ess=675.366681030
4.52 secs, 1492 evals, 1382 steps, improv/step: 0.216 (last = 0.1065), fitn
ess=675.366681030
5.02 secs, 1660 evals, 1550 steps, improv/step: 0.212 (last = 0.1726), fitn
ess=639.065188214
5.52 secs, 1828 evals, 1718 steps, improv/step: 0.205 (last = 0.1488), fitn
ess=428.636072050
6.02 secs, 1998 evals, 1888 steps, improv/step: 0.200 (last = 0.1471), fitn
ess=428.636072050
6.53 secs, 2168 evals, 2058 steps, improv/step: 0.194 (last = 0.1235), fitn
ess=428.636072050
7.03 secs, 2338 evals, 2228 steps, improv/step: 0.189 (last = 0.1294), fitn
ess=421.467654570
7.53 secs, 2506 evals, 2396 steps, improv/step: 0.188 (last = 0.1786), fitn
ess=421.467654570
8.03 secs, 2677 evals, 2567 steps, improv/step: 0.184 (last = 0.1228), fitn
ess=116.717098757
8.54 secs, 2848 evals, 2738 steps, improv/step: 0.178 (last = 0.0877), fitn
ess=116.717098757
9.04 secs, 3018 evals, 2908 steps, improv/step: 0.172 (last = 0.0706), fitn
ess=116.717098757
9.55 secs, 3189 evals, 3079 steps, improv/step: 0.169 (last = 0.1228), fitn
ess=116.717098757
10.06 secs, 3360 evals, 3250 steps, improv/step: 0.168 (last = 0.1462), fit
ness=77.559939475
10.56 secs, 3530 evals, 3420 steps, improv/step: 0.167 (last = 0.1471), fit
ness=55.775174323
11.06 secs, 3697 evals, 3587 steps, improv/step: 0.164 (last = 0.1078), fit
ness=55.775174323
11.56 secs, 3866 evals, 3756 steps, improv/step: 0.162 (last = 0.1302), fit
ness=55.775174323
12.06 secs, 4034 evals, 3924 steps, improv/step: 0.162 (last = 0.1488), fit
ness=53.093054532

Optimization stopped after 4001 steps and 12.293807983398438 seconds
Termination reason: Max number of steps (4000) reached
Steps per second = 325.44838876635714
Function evals per second = 334.3959825589838
Improvements/step = 0.16175
Total function evaluations = 4111

Best candidate found: [0.992043, 1.10023, 0.0962732, 0.522649]

Fitness: 53.093054532

opt = Opt(:GN_ORIG_DIRECT_L, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

24.228527 seconds (98.95 M allocations: 8.149 GiB, 8.92% gc time)
(81.06091854762182, [1.11111, 1.11111, 0.100594, 0.576132], :XTOL_REACHED)

opt = Opt(:GN_CRS2_LM, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 20000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

26.060157 seconds (106.15 M allocations: 8.742 GiB, 8.94% gc time)
(8.128198024607207e-19, [0.7, 0.8, 0.08, 0.5], :XTOL_REACHED)

opt = Opt(:GN_ISRES, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 50000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

167.359077 seconds (607.55 M allocations: 50.034 GiB, 8.80% gc time)
(1.0967650982601886e-13, [0.7, 0.8, 0.08, 0.5], :MAXEVAL_REACHED)

opt = Opt(:GN_ESCH, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[5.0,5.0,5.0,5.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 20000)
@time (minf,minx,ret) = NLopt.optimize(opt,glo_init)

68.752776 seconds (243.02 M allocations: 20.014 GiB, 8.81% gc time)
(720.4749823980964, [2.50782, 2.51209, 0.107538, 0.494919], :MAXEVAL_REACHE
D)

opt = Opt(:LN_BOBYQA, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

1.384155 seconds (4.78 M allocations: 402.706 MiB, 8.50% gc time)
(8.090325025142272e-19, [0.7, 0.8, 0.08, 0.5], :XTOL_REACHED)

opt = Opt(:LN_NELDERMEAD, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-9)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

1.816141 seconds (6.49 M allocations: 547.189 MiB, 8.65% gc time)
(3160.405522281567, [1.0, 1.0, 1.0, 0.865699], :XTOL_REACHED)

opt = Opt(:LD_SLSQP, 4)
lower_bounds!(opt,[0.0,0.0,0.0,0.0])
upper_bounds!(opt,[1.0,1.0,1.0,1.0])
min_objective!(opt, obj.cost_function2)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,loc_init)

0.908700 seconds (3.10 M allocations: 261.275 MiB, 9.01% gc time)
(3160.4056697694823, [1.0, 1.0, 1.0, 0.865765], :XTOL_REACHED)


As expected from other problems the longer sample proves to be extremely challenging for some of the global optimizers. A few give the accurate values, while others seem to struggle with accuracy a lot.

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short)
lower = [0,0,0,0]
upper = [1,1,1,1]
splits = ([0,0.3,0.7],[0,0.3,0.7],[0,0.3,0.7],[0,0.3,0.7])
@time root, x0 = analyze(obj_short,splits,lower,upper)

11.622876 seconds (30.31 M allocations: 1.591 GiB, 9.57% gc time)
(BoxRoot@[NaN, NaN, NaN, NaN], [0.3, 0.3, 0.3, 0.3])

minimum(root)

Box0.0006293336660765138@[0.595183, 0.480868, 0.0638946, 0.500513]

obj = build_loss_objective(prob,Vern9(),L2Loss(t,data),tstops=t,reltol=1e-9,abstol=1e-9)
lower = [0,0,0,0]
upper = [5,5,5,5]
splits = ([0,0.5,1],[0,0.5,1],[0,0.5,1],[0,0.5,1])
@time root, x0 = analyze(obj_short,splits,lower,upper)

0.221035 seconds (927.43 k allocations: 80.229 MiB, 16.21% gc time)
(BoxRoot@[NaN, NaN, NaN, NaN], [0.5, 0.5, 0.5, 0.5])

minimum(root)

Box0.012650644127007372@[0.166695, 0.999023, 0.389981, 0.499926]


# Conclusion

It is observed that lower tolerance lead to higher accuracy but too low tolerance could affect the convergance time drastically. Also fitting a shorter timespan seems to be easier in comparision (quite intutively). NLOpt methods seem to give great accuracy in the shorter problem with a lot of the algorithms giving 0 fitness, BBO performs very well on it with marginal change with tol values. In case of global optimization of the longer problem there is some difference in the perfomance amongst the algorithms with :LNBOBYQA giving accurate results for the local optimization and :GNISRES :GNCRS2LM in case of the global give the highest accuracy. BBO also fails to perform too well in the case of the longer problem. QuadDIRECT performs well in case of the shorter problem but fails to give good results in the longer version.