Pytorch code for a parameter estimation Neural Network applied to the SIR mathematical model of epidemiology.
Consider the SIR model
In other words, this neural network determines the infectivity of a disease from the infected data (Grimm et al, 2021).
The default data, from which the Neural Network (NN) determines the parameter, is the solution of the SIR equations. The neural network does not have any information of the real value of the infectivity parameter
The data is transformed using a log function followed by a min-max transformation such that the input data of the NN is comprised between 0 and 1. The code uses a neural network with a SELU activation function and its architecture is composed by 10 layers of 128 neurons each and the loss functions are the MSE of the time series and the residuals from the SIR differential equations. A weighting/scaling procedure is in place to ensure the proper evolution of the NN (Grimm et al, 2021).
Grimm, V. et al. (2021) ‘Estimating the time-dependent contact rate of SIR and SEIR models in mathematical epidemiology using physics-informed neural networks’, ETNA - Electronic Transactions on Numerical Analysis, 56, pp. 1–27. Available at: https://doi.org/10.1553/etna_vol56s1.