Coverage for aftercovid/models/_base_sir_estimation.py : 100%
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"""
3Common methods about training, predicting for :epkg:`SIR` models.
4"""
5import pprint
6import numpy
7from sympy import Symbol, diff as sympy_diff
8from ..optim import SGDOptimizer
11class BaseSIREstimation:
12 """
13 Common methods about training, predicting for :epkg:`SIR` models.
14 """
16 def _check_fit_predict(self, X, y=None):
17 if not isinstance(X, numpy.ndarray):
18 raise TypeError("X must be a numpy array.")
19 if len(X.shape) != 2:
20 raise ValueError("X must be a matrix.")
21 clq = self.quantity_names
22 if len(clq) != X.shape[1]:
23 raise ValueError(
24 "Unexpected number of columns, got {}, expected {}.".format(
25 X.shape[1], len(clq)))
26 sums = numpy.sum(X, axis=1)
27 mi, ma = sums.min(), sums.max()
28 df = abs(ma - mi)
29 if df > abs(ma) * 1e-3:
30 raise ValueError(
31 "All rows must sum up to the same amount. Current "
32 "range: [{}, {}].".format(mi, max))
33 if y is not None:
34 if not isinstance(y, numpy.ndarray):
35 raise TypeError("y must be a numpy array.")
36 if y.shape != X.shape:
37 raise ValueError(
38 "Unexpected shape of y, got {}, expected {}.".format(
39 y.shape, X.shape))
40 return clq, ma
42 def fit(self, X, y, t=0, max_iter=100,
43 learning_rate_init=0.1, lr_schedule='constant',
44 momentum=0.9, power_t=0.5, early_th=None,
45 min_threshold=None, max_threshold=None,
46 verbose=False):
47 """
48 Fits a model :class:`BaseSIR <aftercovid.models._base_sir.BaseSIR>`.
50 :param X: known values for every quantity at time *t*,
51 every column is mapped to the list returned by
52 :meth:`quantity_names <aftercovid.models._base_sir.quantity_names>`
53 :param y: known derivative for every quantity at time *t*,
54 comes in the same order as *X*,
55 both *X* and *y* have the same shape.
56 :param t: implicit feature
57 :param max_iter: number of iteration
58 :param learning_rate_init: see :class:`SGDOptimizer
59 <aftercovid.optim.SGDOptimizer>`
60 :param lr_schedule: see :class:`SGDOptimizer
61 <aftercovid.optim.SGDOptimizer>`
62 :param momentum: see :class:`SGDOptimizer
63 <aftercovid.optim.SGDOptimizer>`
64 :param power_t: see :class:`SGDOptimizer
65 <aftercovid.optim.SGDOptimizer>`
66 :param early_th: see :class:`SGDOptimizer
67 <aftercovid.optim.SGDOptimizer>`
68 :param verbose: see :class:`SGDOptimizer
69 <aftercovid.optim.SGDOptimizer>`
70 :param min_threshold: see :class:`SGDOptimizer
71 <aftercovid.optim.SGDOptimizer>`
72 :param max_threshold: see :class:`SGDOptimizer
73 <aftercovid.optim.SGDOptimizer>`
75 The training needs two steps. The first one creates a training
76 datasets. The second one estimates the coefficients by using
77 a stochastic gradient descent (see :class:`SGDOptimizer
78 <aftercovid.optim.SGDOptimizer>`).
79 Let's use a SIR model (see :class:`CovidSIR
80 <aftercovid.models.CovidSIR>`).as an example.
81 Let's denote the parameters as :math:`\\Omega`
82 and :math:`Z_1=S`, ..., :math:`Z_4=R`.
83 The model is defined by
84 :math:`\\frac{dZ_i}{dt} = f_i(\\Omega, Z)`
85 where :math:`Z=(Z_1, ..., Z_4)`.
86 *y* is used to compute the expected derivates
87 :math:`\\frac{dZ_i}{dt}`. The loss function is defined as:
89 .. math::
91 L(\\Omega,Z) = \\sum_{i=1}^4 \\left( f_i(\\Omega,Z) -
92 \\frac{dZ_i}{dt}\\right)^2
94 Then the gradient is:
96 .. math::
98 \\frac{\\partial L(\\Omega,Z)}{\\partial\\Omega} =
99 2 \\sum_{i=1}^4 \\frac{\\partial f_i(\\Omega,Z)}{\\partial\\Omega}
100 \\left( f_i(\\Omega,Z) - \\frac{dZ_i}{dt} \\right)
102 A stochastic gradient descent takes care of the rest.
103 """
104 self._fit(X, y, t=0, max_iter=max_iter,
105 learning_rate_init=learning_rate_init,
106 lr_schedule=lr_schedule, momentum=momentum,
107 power_t=power_t, verbose=verbose, early_th=early_th,
108 min_threshold=min_threshold, max_threshold=max_threshold)
109 return self
111 def _losses_sympy(self):
112 """
113 Builds the loss functions using :epkg:`sympy`,
114 one for each quantity.
115 """
116 clq = self.quantity_names
117 res = []
118 for name in clq:
119 sym = Symbol('d' + name)
120 eq = self._eq[name]
121 lo = (eq - sym) ** 2
122 la = self._lambdify_('loss-%s' % name, lo, derivative=True)
123 res.append((lo, la))
124 return res
126 def _grads_sympy(self):
127 """
128 Builds the gradient for every parameter,
129 exactly number of quantities multiplied by
130 the number of parameters.
131 """
132 losses = self._losses_sympy()
133 clq = self.quantity_names
134 prn = self.param_names
135 res = []
136 for name, eq in zip(clq, losses):
137 row = []
138 for pn in prn:
139 pns = Symbol(pn)
140 df = sympy_diff(eq[0], pns)
141 gr = self._lambdify_('grad-%s/%s' % (name, pn), df,
142 derivative=True)
143 row.append((df, gr))
144 res.append(row)
145 return res
147 def predict(self, X, t=0):
148 """
149 Predicts the derivative at time *t*.
151 :param X: known values for every quantity at time *t*,
152 every column is mapped to the list returned by
153 :meth:`quantity_names <aftercovid.models._base_sir.quantity_names>`
154 :param t: implicit feature
155 :return: predictive derivative
156 """
157 N = self._check_fit_predict(X)[1]
158 err = abs(N - self['N']) / N
159 if err > 1e-4:
160 raise ValueError( # pragma: no cover
161 "All rows must sum up to {} not {}.".format(self['N'], N))
162 n = X.shape[0]
163 C = X.shape[1]
164 pred = numpy.empty(X.shape, dtype=X.dtype)
165 x = self.vect(t=t)
166 for i in range(t, t + n):
167 x[-1] = i
168 x[:C] = X[i, :]
169 for c, f in enumerate(self._leqa):
170 pred[i, c] = f(*x)
171 return pred
173 def score(self, X, y, t=0):
174 """
175 Scores the predictions. Returns L2 norm
176 divided by the number of rows.
178 :param X: known values for every quantity at time *t*,
179 every column is mapped to the list returned by
180 :meth:`quantity_names <aftercovid.models._base_sir.quantity_names>`
181 :param y: expected values
182 :param t: implicit feature
183 :return: predictive derivative
184 """
185 self._check_fit_predict(X, y)
186 pred = self.predict(X, t=t)
187 delta = (pred - y) ** 2
188 return numpy.sum(delta) / X.shape[0]
190 def score_l1(self, X, y, t=0):
191 """
192 Scores the predictions. Returns L1 norm
193 divided by the number of rows and the population.
195 :param X: known values for every quantity at time *t*,
196 every column is mapped to the list returned by
197 :meth:`quantity_names <aftercovid.models._base_sir.quantity_names>`
198 :param y: expected values
199 :param t: implicit feature
200 :return: predictive derivative
201 """
202 self._check_fit_predict(X, y)
203 pred = self.predict(X, t=t)
204 delta = numpy.abs(pred - y)
205 return numpy.sum(delta) / (X.shape[0] * self['N'])
207 def _fit(self, X, y, t, max_iter,
208 learning_rate_init, lr_schedule,
209 momentum, power_t, early_th, min_threshold,
210 max_threshold, verbose):
211 '''
212 See method :meth:`fit
213 <aftercovid.models._base_sir_estimation.BaseSIREstimation.fit>`
214 and :class:`SGDOptimizer <aftercovid.optim.SGDOptimizer>`.
215 '''
216 N = self._check_fit_predict(X, y)[1]
217 err = abs(N - self['N']) / N
218 if err > 1e-4:
219 raise ValueError( # pragma: no cover
220 "All rows must sum up to {} not {}.".format(self['N'], N))
222 # loss and gradients functions
223 losses = self._losses_sympy()
224 grads = self._grads_sympy()
225 xy0 = self.vect(t=0, derivative=True)
227 def pformat(d): # pragma: no cover
228 nd = {str(k): (str(v), type(v), type(k)) for k, v in d.items()}
229 return pprint.pformat(nd)
231 def fct_loss(coef, X, y):
232 'Computes the loss function for every X and y.'
233 xy0[self._val_ind[1]:self._val_ind[2]] = coef
235 res = 0.
236 for i in range(X.shape[0]):
237 xy0[-1] = i
238 xy0[:X.shape[1]] = X[i, :]
239 xy0[-self._val_ind[1]:] = y[i, :]
240 for loss in losses:
241 res += loss[1](*xy0)
242 return res
244 def fct_grad(coef, x, y, i):
245 'Computes the gradient function for every X and y.'
246 xy0[:self._val_ind[1]] = x
247 xy0[self._val_ind[1]:self._val_ind[2]] = coef
248 xy0[-self._val_ind[1]:] = y
250 res = numpy.zeros((coef.shape[0], ), dtype=x.dtype)
251 for row in grads:
252 for i, g in enumerate(row):
253 res[i] += g[1](*xy0)
254 return res
256 pnames = self.param_names
257 coef = numpy.array([self[p] for p in pnames])
258 lrn = 1. / (N ** 1.5)
259 sgd = SGDOptimizer(
260 coef, learning_rate_init=learning_rate_init * lrn,
261 lr_schedule=lr_schedule, momentum=momentum,
262 power_t=power_t, min_threshold=min_threshold,
263 max_threshold=max_threshold)
265 sgd.train(X, y, fct_loss, fct_grad, max_iter=max_iter,
266 early_th=early_th, verbose=verbose)
268 # uses trained coefficients
269 coef = sgd.coef
270 for n, c in zip(pnames, coef):
271 self[n] = c
272 self.iter_ = sgd.iter_