Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# coding: utf-8 

2""" 

3Implementation of a model for epidemics propagation. 

4""" 

5import numpy 

6from sklearn.base import BaseEstimator, RegressorMixin 

7from .covid_sird import CovidSIRD 

8from .covid_sird_cst import CovidSIRDc 

9 

10 

11class EpidemicRegressor(BaseEstimator, RegressorMixin): 

12 """ 

13 Follows :epkg:`scikit-learn` API. 

14 Trains a model on observed data from an epidemic. 

15 

16 :param model: model to train, `'SIR'` or `'SIRD'` 

17 refers to `CovidSIRD <aftercovid.models.CovidSIRD>`, 

18 `SIRDc` refers to `CovidSIRDc 

19 <aftercovid.models.CovidSIRDc>` 

20 :param t: implicit feature 

21 :param max_iter: number of iteration 

22 :param learning_rate_init: see :class:`SGDOptimizer 

23 <aftercovid.optim.SGDOptimizer>` 

24 :param lr_schedule: see :class:`SGDOptimizer 

25 <aftercovid.optim.SGDOptimizer>` 

26 :param momentum: see :class:`SGDOptimizer 

27 <aftercovid.optim.SGDOptimizer>` 

28 :param power_t: see :class:`SGDOptimizer 

29 <aftercovid.optim.SGDOptimizer>` 

30 :param early_th: see :class:`SGDOptimizer 

31 <aftercovid.optim.SGDOptimizer>` 

32 :param verbose: see :class:`SGDOptimizer 

33 <aftercovid.optim.SGDOptimizer>` 

34 :param min_threshold: see :class:`SGDOptimizer 

35 <aftercovid.optim.SGDOptimizer>`, if `'auto'`, 

36 the value depends on the models, if is `0.01` 

37 for model `SIR`, it means every coefficient must 

38 be greater than 0.01. 

39 :param max_threshold: see :class:`SGDOptimizer 

40 <aftercovid.optim.SGDOptimizer>`, upper bound 

41 :param init: dictionary, initializes the model 

42 with this parameters 

43 

44 Once trained the model holds a member `model_` 

45 which contains the trained model and `iter_` 

46 which holds the number of training iteration. 

47 It also keep track of the coefficients in a dictionary 

48 in attribute `coef_`. 

49 """ 

50 

51 def __init__(self, model='SIR', t=0, max_iter=100, 

52 learning_rate_init=0.1, lr_schedule='constant', 

53 momentum=0.9, power_t=0.5, early_th=None, 

54 min_threshold='auto', max_threshold='auto', 

55 verbose=False, init=None): 

56 if init is not None: 

57 if isinstance(init, EpidemicRegressor): 

58 if hasattr(init, 'coef_'): 

59 init = init.coef_.copy() 

60 else: 

61 init = None # pragma: no cover 

62 elif not isinstance(init, dict): 

63 raise TypeError( 

64 "init must be a dictionary not {}.".format(type(init))) 

65 BaseEstimator.__init__(self) 

66 RegressorMixin.__init__(self) 

67 self.t = t 

68 self.model = model 

69 self.max_iter = max_iter 

70 self.learning_rate_init = learning_rate_init 

71 self.lr_schedule = lr_schedule 

72 self.momentum = momentum 

73 self.power_t = power_t 

74 self.early_th = early_th 

75 self.verbose = verbose 

76 if min_threshold == 'auto': 

77 if model.upper() in ('SIR', 'SIRD'): 

78 min_threshold = 0.0001 

79 elif model.upper() in ('SIRC', ): 

80 pmin = dict(beta=0.001, nu=0.0001, mu=0.0001, 

81 a=-1., b=0., c=0.) 

82 min_threshold = numpy.array( 

83 [pmin[k[0]] for k in CovidSIRDc.P0]) 

84 elif model.upper() in ('SIRDC'): 

85 pmin = dict(beta=0.001, nu=0.001, mu=0.001, 

86 a=-1., b=0., c=0.) 

87 min_threshold = numpy.array( 

88 [pmin[k[0]] for k in CovidSIRDc.P0]) 

89 if max_threshold == 'auto': 

90 if model.upper() in ('SIR', 'SIRD'): 

91 max_threshold = 1. 

92 elif model.upper() in ('SIRC', 'SIRDC'): 

93 pmax = dict(beta=1., nu=0.5, mu=0.5, 

94 a=0., b=4., c=2.) 

95 max_threshold = numpy.array( 

96 [pmax[k[0]] for k in CovidSIRDc.P0]) 

97 self.min_threshold = min_threshold 

98 self.max_threshold = max_threshold 

99 self._get_model() 

100 self.init = init 

101 if init is not None: 

102 self.coef_ = init 

103 

104 def _get_model(self): 

105 if self.model.lower() in ('sir', 'sird'): 

106 return CovidSIRD() 

107 if self.model.lower() in ('sirc', 'sirdc'): 

108 return CovidSIRDc() 

109 raise ValueError( 

110 "Unknown model name '{}'.".format(self.model)) 

111 

112 def fit(self, X, y): 

113 """ 

114 Trains a model to approximate its derivative as much as 

115 possible. 

116 """ 

117 if not hasattr(self, 'model_'): 

118 self.model_ = self._get_model() 

119 self.model_.rnd() 

120 total = numpy.sum(X, axis=1) 

121 mi, ma = total.min(), total.max() 

122 err = (ma - mi) / mi 

123 if err > 1e-5: 

124 raise RuntimeError( # pragma: no cover 

125 "Population is not constant, in [{}, {}].".format( 

126 mi, ma)) 

127 if self.init is not None: 

128 for k, v in self.init.items(): 

129 self.model_[k] = v 

130 self.model_['N'] = (ma + mi) / 2 

131 self.model_.fit( 

132 X, y, learning_rate_init=self.learning_rate_init, 

133 max_iter=self.max_iter, early_th=self.early_th, 

134 verbose=self.verbose, lr_schedule=self.lr_schedule, 

135 power_t=self.power_t, momentum=self.momentum, 

136 min_threshold=self.min_threshold, 

137 max_threshold=self.max_threshold) 

138 self.iter_ = self.model_.iter_ 

139 self.coef_ = self.model_.params_dict 

140 return self 

141 

142 def predict(self, X): 

143 """ 

144 Predicts the derivatives. 

145 """ 

146 if not hasattr(self, 'model_'): 

147 raise RuntimeError("Model was not trained.") 

148 return self.model_.predict(X) 

149 

150 def simulate(self, X, n=7): 

151 """ 

152 Predicts and simulates the epidemics. 

153 Every row of *X* is a starting point, 

154 the function then simulates the epidemics for the next 

155 *n* days for every starting point. 

156 

157 :param X: data 

158 :param n: number of days 

159 :return: quantities, matrix of shape 

160 *(X.shape[0], n, number of parameters)* 

161 """ 

162 if not hasattr(self, "model_"): 

163 raise RuntimeError( # pragma: no cover 

164 "Model was not trained.") 

165 clq = self.model_.quantity_names 

166 if len(clq) != X.shape[1]: 

167 raise RuntimeError( # pragma: no cover 

168 "Unapexected shape for X ({}), expecting {} columns." 

169 "".format(X.shape, len(clq))) 

170 res = None 

171 for i in range(X.shape[0]): 

172 for k, v in zip(clq, X[i]): 

173 self.model_[k] = v 

174 pred = self.model_.iterate2array(n=n, derivatives=False) 

175 if res is None: 

176 res = numpy.zeros((X.shape[0], ) + pred.shape) 

177 res[i, :, :] = pred 

178 return res 

179 

180 def predict_many(self, X, n=7): 

181 """ 

182 Predicts the derivatives and the series 

183 for many days. 

184 

185 :param X: series 

186 :param n: number of days 

187 :return: derivates and series, return shape is 

188 *(X.shape[0], number of parameters, n)* 

189 """ 

190 if not hasattr(self, 'model_'): 

191 raise RuntimeError("Model was not trained.") # pragma: no cover 

192 deri = numpy.empty(X.shape + (n, )) 

193 curv = numpy.empty(X.shape + (n, )) 

194 for i in range(0, n): 

195 d = self.predict(X) 

196 deri[:, :, i] = d 

197 X += d 

198 curv[:, :, i] = X 

199 return deri, curv 

200 

201 def score(self, X, y, norm=None): 

202 """ 

203 Scores the prediction of the derivatives. 

204 

205 :param X: data 

206 :param y: expected derivatives 

207 :param norm: norm to return the norm used to optimize (L2) 

208 or 'L1' to return the L1 norm 

209 :return: score 

210 """ 

211 if not hasattr(self, 'model_'): 

212 raise RuntimeError( # pragma: no cover 

213 "Model was not trained.") 

214 if norm is None: 

215 return self.model_.score(X, y) 

216 if norm.lower() == 'l1': 

217 return self.model_.score_l1(X, y) 

218 raise ValueError("Unexpected norm %r." % norm)