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
11class EpidemicRegressor(BaseEstimator, RegressorMixin):
12 """
13 Follows :epkg:`scikit-learn` API.
14 Trains a model on observed data from an epidemic.
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
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 """
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
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))
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
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)
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.
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
180 def predict_many(self, X, n=7):
181 """
182 Predicts the derivatives and the series
183 for many days.
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
201 def score(self, X, y, norm=None):
202 """
203 Scores the prediction of the derivatives.
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)