(*---------------------------------------------------------------------*)
(* The actual run in the paper for tolerance 10^-9 *)
(*---------------------------------------------------------------------*)
(* Logistic Growth ODE *)
(* RKT7(5)sa *)
Needs["DifferentialEquations`NDSolveProblems`"];
New75cvec = {2509146235436/5087746235499, 1350370294159/3408657898282,
    3252443538291/4522619808934, 235723152996/10608700391405, 
   310760205251/400918996038, 1/2, 1, 1};
New75bvec = {-878661591969/2977574671790, 113396230556/2709258533071, 
   1463959570440/2917571478191, 222546947234/4419440739125, 
   3491876531739/7346997574445, 
   649395555859/2043410174806, -311611088581/1961303128474, 
   256166179812/3828598790879, 0};
New75evec = {-878661591969/2977574671790, 113396230556/2709258533071, 
    1463959570440/2917571478191, 222546947234/4419440739125, 
    3491876531739/7346997574445, 
    649395555859/2043410174806, -311611088581/1961303128474, 
    256166179812/3828598790879, 0} - {-522012459733/1481211415688, 
    104797793945/2255020572942, 1585964307576/3483995656921, 
    339661082807/4016940902768, 868681294941/1602978375592, 
    788400671307/2849252062316, -765424471000/6187037101927, 
    61731523173/2902151573873, 1/20};
New75amat = {{2509146235436/5087746235499}, {617954070027/
     2562646188400, 
    638714441812/4120209690065}, {-301533143356/
     3409145176049, -2631922871921/5416230055909, 
    2739512869868/2117855932099}, {41413341307/3389709942745, 
    15718215260/3297610944251, -7568154684/1622901758375, 
    16276114867/1644180493967}, {-297182739254/1683969669677, 
    582170851839/2035654146344, 76215464473/20514042277991, 
    752399359454/3719737864487, 
    837521133301/1822192981660}, {5260131130959/2695616599657, 
    3403289853008/6000854915727, -645218926486/1328789175843, 
    319229450978/2340066684127, -11878038161387/
     7344669137326, -129277603880/2480564774447}, {3213873918485/
     3313536732549, -1045710975791/3635704707653, 
    9305863180976/4984307714753, -1459922494763/
     3239497629219, -3815867383322/3490900707541, 
    2736079125313/2960105948232, -3427392396131/
     3685775496659}, {-878661591969/2977574671790, 
    113396230556/2709258533071, 1463959570440/2917571478191, 
    222546947234/4419440739125, 3491876531739/7346997574445, 
    649395555859/2043410174806, -311611088581/1961303128474, 
    256166179812/3828598790879}};
New75Coefficients[7, p_] := 
  N[{New75amat, New75bvec, New75cvec, New75evec}, p];

(*counting time*)
(*Logistic Growth ODE*)
logisticODE[y_, r_, K_] := r y (1 - y/K)

(*True parameters for synthetic data*)
rTrue = 2/3;
KTrue = 20;
y0 = 4;

(*Time points and synthetic data*)
times = N@Range[0, 10, 1/2];
trueSol = 
  y /. NDSolve[{y'[t] == logisticODE[y[t], rTrue, KTrue], y[0] == y0},
      y, {t, 0, 10}, WorkingPrecision -> 33, AccuracyGoal -> 24, 
     PrecisionGoal -> 24][[1]];
data = SetAccuracy[Table[{t, trueSol[t]}, {t, times}], 24];

(*Objective loss function:squared error between model and data*)
computeLoss[r_, K_] := 
 Module[{sol, pred}, 
  sol = NDSolveValue[{y'[t] == r y[t] (1 - y[t]/K), y[0] == y0}, 
    y, {t, 0, 10}, 
    Method -> {"ExplicitRungeKutta", "DifferenceOrder" -> 7, "EmbeddedDifferenceOrder" -> 5,
      "Coefficients" -> New75Coefficients}, PrecisionGoal -> 9, 
    AccuracyGoal -> 9];
  pred = Table[sol[t], {t, times}];
  Total[(pred - data[[All, 2]])^2]]

(*Gradient approximation using central differences*)
computeGradient[r_, K_, \[Delta]_ : 0.000005] := 
 Module[{dLdr, dLdK}, 
  dLdr = (computeLoss[r + \[Delta], K] - 
      computeLoss[r - \[Delta], K])/(2 \[Delta]);
  dLdK = (computeLoss[r, K + \[Delta]] - 
      computeLoss[r, K - \[Delta]])/(2 \[Delta]);
  {dLdr, dLdK}]

(*Momentum-based Gradient Descent*)
ClearAll[trainModel]
trainModel[rInit_, KInit_, \[Eta]_ : 0.0001, \[Gamma]_ : 0.9, 
  maxIter_ : 1400] := 
 Module[{r = rInit, K = KInit, vr = 0, vK = 0, lossLog = {}, 
   paramLog = {}}, Do[AppendTo[lossLog, computeLoss[r, K]];
   AppendTo[paramLog, {r, K}];
   If[Mod[i, 200] == 0, 
    Print["Iter ", i, ": Loss = ", lossLog[[-1]], ", r = ", r, 
     ", K = ", K]];
   (*Gradient computation*){gr, gK} = computeGradient[r, K];
   (*Update momentum*)vr = \[Gamma] vr + \[Eta] gr;
   vK = \[Gamma] vK + \[Eta] gK;
   (*Update parameters*)r = r - vr;
   K = K - vK;
   (*Enforce bounds*)r = Clip[r, {0.1, 2.0}];
   K = Clip[K, {10.0, 100.0}];, {i, 1, maxIter}];
  {r, K, lossLog, paramLog}]

(*Run the optimizer*)
{runTime, {rEst, KEst, losses, params}} = 
  AbsoluteTiming[trainModel[6/10, 25]];
Print["Execution Time (seconds): ", runTime];

(*Relative Errors*)
epsR = Abs[(rEst - rTrue)/rTrue]*1;
epsK = Abs[(KEst - KTrue)/KTrue]*1;
Print["Relative Error in r (%): ", epsR];
Print["Relative Error in K (%): ", epsK];