'NonlinearModelFit' for SIR ODE model, and adding a penalty function (soft constraint)

Question

This is a modification of the question Is it normal that adding constraints on the parameters of NonlinearModelFit increases execution time 100 times??

I want a program that fits a SIR model to reported new infections data. The SIR model works, it produces a vector of 33 components (estimated weekly new infections)

pobTotal = 4950738;
 mu = 7*41157/365 // N(* weekly births *);
d = mu/pobTotal;(*constant population*)
reported = {107, 135, 612, 195, 626, 619, 491, 1164, 1137, 511, 1036, 
   1144, 2650, 3162, 6074, 6693, 8253, 6639, 6148, 4345, 3141, 1958, 
   1130, 484, 356, 296, 195, 121, 208, 101, 67, 128, 20};
data = Thread[{Range[1, 33], reported}];
T = Length[reported ];n=7;(*discretization to days, to avoid long times taken by NIntegrate*)
SIR[\[Gamma]_?NumericQ, \[Beta]_?NumericQ, s0_?NumericQ, i0_?NumericQ] :=
  Module[{ssol, isol, RiemSum, fd},
  {ssol, isol} = NDSolveValue[{
     s'[t] == mu - (\[Beta]*s[t]*i[t]/pobTotal) - d*s[t],
     i'[t] == (\[Beta]*s[t]*i[t]/pobTotal) - (\[Gamma] + d)*i[t],
     s[0] == s0, i[0] == i0
     },
    {s, i},
    {t, 0, T},
    MaxStepSize -> 400];
  fd = \[Beta]*ssol[Range[0, n T]/n]*isol[Range[0, n T]/n]/pobTotal;
  RiemSum = Total[Partition[Drop[fd, 1], n], {2}]/n]
  SIR[1,1,10^6,22]
  ListLinePlot[SIR[1,1,10^6,22]]

I tried 'NonlinearModelFit' which works very quickly, but produces negative parameters, and the plot also looks nonsense.

Timing[nlm = NonlinearModelFit[
    data,  p*SIR[\[Gamma], \[Beta], s0, i0], {{\[Gamma],1}, {\[Beta],1.8}, {s0,4254150}, {i0,146},{p,.023}}, t]] // Quiet
    {\[Gamma],\[Beta],s0,i0,p}/. nlm["BestFitParameters"]
    ListPlot[nlm[t]]

It seems the input to 'NonlinearModelFit' should be a function, not a vector, which leads to the nonsensical output. This suggests I should use some other fitting command.

A second question: I must also add constraints on the parameters. These multiply the execution times by hundreds, and it would be nice to be able to replace constraints by a penalty function (soft constraint). Is it possible to modify the objective of 'NonlinearModelFit' by a penalty function ? Or are we forced to use 'NMinimize', and give up all the statistics diagnostics of 'NonlinearModelFit'?

Answer 1

The SIR function (specifically p SIR[γ, β, s0, i0]) now produces a list of predicted values. However, NonlinearModelFit only allows models that produce a single value for each observation.

To get around this obstacle one can find the values of the parameters that minimize the sum of squares or the root mean square error (rmse) which is equivalent in this case to finding the maximum likelihood estimates.

rmse[γ_?NumericQ, β_?NumericQ, s0_?NumericQ, i0_?NumericQ, p_?NumericQ] := Module[{z},
  z = p*SIR[γ, β, s0, i0];
  Sqrt[Total[(data[[All, 2]] - z)^2]/(Length[data] - 5)]]

sol = FindMinimum[{rmse[γ, β, s0, i0, p], γ > 0 && β > 0 && s0 > 0 && i0 > 0 && p > 0},
  {{γ, 1}, {β, 1.8}, {s0, 5254150}, {i0, 146}, {p, 0.023}}]

(* {417.92, {γ -> 0.00356632, β -> 0.670843, s0 -> 4.1135*10^6, i0 -> 385.799, p -> 0.0135433}} *)

(* Construct interpolation function so that predictions can be made from non-integer values *)
data2 = Transpose[{Range[33], p SIR[γ, β, s0, i0] /. sol[[2]]}];
g = Interpolation[data2]

Show[ListPlot[data], Plot[g[t], {t, 1, 33}]]