Coverage for src/ensae_teaching_cs/special/propagation_epidemic.py : 85%
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"""
2@file
3@brief Simple simulation of an epidemic. It makes a couple of assumption
4on how the disease is transmitted and its effect on a person. The user
5can specify many parameters.
6"""
7import random
8import copy
9import os
10from collections import Counter
11from pyquickhelper.loghelper import noLOG
12from ..helpers.pygame_helper import wait_event, empty_main_loop
15class Point:
16 """
17 Defines a point.
18 """
20 def __init__(self, x, y):
21 """
22 constructor
24 @param x x
25 @param y y
26 """
27 self.x, self.y = float(x), float(y)
29 def norm(self):
30 """
31 return the norm l2
32 """
33 return (self.x ** 2 + self.y ** 2) ** 0.5
35 def __add__(self, p):
36 """
37 addition
38 """
39 return Point(self.x + p.x, self.y + p.y)
41 def __sub__(self, p):
42 """
43 soustraction
44 """
45 return Point(self.x - p.x, self.y - p.y)
47 def __mul__(self, p):
48 """
49 multiplication by a scalar
50 """
51 return Point(self.x * p, self.y * p)
53 def __div__(self, p):
54 """
55 division by a scalar
56 """
57 return Point(self.x / p, self.y / p)
59 def __iadd__(self, p):
60 """
61 addition inplace
62 """
63 self.x += p.x
64 self.y += p.y
65 return self
67 def __isub__(self, p):
68 """
69 soustraction inplace
70 """
71 self.x -= p.x
72 self.y -= p.y
73 return self
75 def __imul__(self, p):
76 """
77 multiplication by a scalar inplace
78 """
79 self.x *= p
80 self.y *= p
81 return self
83 def __idiv__(self, p):
84 """
85 divsion by a scalar inplace
86 """
87 self.x /= p
88 self.y /= p
89 return self
92class Rect:
93 """
94 Defines a rectangle.
95 """
97 def __init__(self, a, b):
98 """
99 constructor
101 @param a Point
102 @param b Point
103 """
104 self.a = a
105 self.b = b
107 def limit(self, pos):
108 """
109 tells if point *pos* belongs to the area
110 defined by the rectangle
111 """
112 r = False
113 if pos.x < self.a.x:
114 pos.x = self.a.x
115 r = True
116 if pos.y < self.a.y:
117 pos.y = self.a.y
118 r = True
119 if pos.x > self.b.x:
120 pos.x = self.b.x
121 r = True
122 if pos.y > self.b.x:
123 pos.y = self.b.y
124 r = True
125 return r
128class Person:
129 """
130 defines a person for the simulation
132 colors
134 * 0: sain
135 * 1: malade
136 * 2: mort
137 * 3: gueris
139 A person moves by drawing a random gaussian vector added to
140 its current acceleration.
141 """
142 colors = {0: (255, 255, 255), 1: (0, 255, 0),
143 2: (0, 0, 0), 3: (50, 50, 200)}
145 def __init__(self, position, borne_pos=None, borne_acc=None,
146 alea_acc=5, sick_vit=0.25, rayon=10, nb_day=70,
147 prob_die=0.5, prob_cont=0.5):
148 """
149 constructor
151 @param position position
152 @param borne_pos upper bound for the position (rectangle)
153 @param borne_acc upper bound for the acceleration (rectangle)
154 @param alea_acc sigma to draw random acceleration
155 @param sick_vit when people are sick, they go slower, this muliplies
156 the acceleration by a factor
157 @param rayon radius, below that distance, a sick person is contagious for the neighbours
158 @param nb_day number of days a person will be sick, after, the person
159 either recovers, either dies
160 @param prob_die probability to die at each iteration
161 @param prob_cont probability to transmit the disease to a neighbour
162 """
163 self.pos = position
164 self.vit = Point(0, 0)
165 self.acc = Point(0, 0)
166 self.state = 0
167 self.alea_acc = alea_acc
168 self.borne_pos = borne_pos
169 self.borne_acc = borne_acc
171 self.sick_vit = sick_vit
172 self.rayon = rayon
173 self.nb_day = nb_day
174 self.prob_die = prob_die
175 self.prob_cont = prob_cont
177 # memorize the day the person got sick
178 self._since = 0
180 def __str__(self):
181 """
182 usual
183 """
184 return str(self.__dict__)
186 def distance(self, p):
187 """
188 return the distance between this person and another one
189 """
190 d = self.pos - p.pos
191 return d.norm()
193 def _get_new_acceleration(self):
194 """
195 update the acceleration by adding a random gaussian vector
196 to the current acceleration, check that acceleration
197 is not beyond some boundary
198 """
199 x = random.gauss(0, self.alea_acc)
200 y = random.gauss(0, self.alea_acc)
201 res = Point(x, y)
202 if self.borne_acc is not None:
203 r = self.borne_acc.limit(res)
204 if r:
205 self.acc = Point(0, 0)
206 self.vit = Point(0, 0)
207 return res
209 def state_evolution(self, population):
210 """
211 update the state of the person: healthy --> sick --> cured or dead
213 @param population sets of other persons
215 The function updates the state of the persons.
216 One of steps involves looking over the entire population to check
217 if some sick people are close enough to transmis the disease.
218 """
219 if self.state in [2, 3]:
220 return
221 elif self.state == 1:
222 if self._since < self.nb_day:
223 self._since += 1
224 else:
225 k = random.random()
226 if k < self.prob_die:
227 self.state = 2
228 else:
229 self.state = 3
230 self._since = 0
231 elif self.state == 0:
232 alls = []
233 for p in population:
234 if p.state != 1:
235 continue
236 d = self.distance(p)
237 if d <= self.rayon:
238 alls.append(p)
240 for k in alls:
241 p = random.random()
242 if p <= self.prob_cont:
243 self.state = 1
244 break
245 else:
246 raise Exception("impossible")
248 def evolution(self, dt, population):
249 """
250 update the population, random acceleration
252 @param dt time delta (only used to update the position)
253 @param population other set of people
255 The function updates the state of the person,
256 draws a new acceleration and updates the position.
257 """
258 self.state_evolution(population)
260 if self.state == 1:
261 dt *= self.sick_vit
262 elif self.state == 2:
263 dt = 0
265 self.pos += self.vit * dt
266 self.vit += self.acc * dt
267 self.acc = self._get_new_acceleration()
268 if self.borne_pos is not None:
269 r = self.borne_pos.limit(self.pos)
270 if r:
271 self.acc = Point(0, 0)
272 self.vit = Point(0, 0)
275class EpidemicPopulation:
276 """
277 defines a population
278 """
280 def __init__(self, cote=500, nb=(100, 1), **params):
281 """
282 constructeur
284 @param cote size of the zone person move
285 @param nb tuple number of people (healthy, sick)
286 @param params others parameters
288 Draws a population.
289 """
290 if cote is None:
291 pass
292 else:
293 self.cote = cote
294 self.gens = []
295 for i in range(0, nb[0]):
296 p = Person(Point(random.randint(0, cote), random.randint(0, cote)),
297 Rect(Point(0, 0), Point(cote, cote)),
298 **params)
299 self.gens.append(p)
300 for i in range(0, nb[1]):
301 p = Person(Point(random.randint(0, cote), random.randint(0, cote)),
302 Rect(Point(0, 0), Point(cote, cote)),
303 **params)
304 p.state = 1
305 self.gens.append(p)
307 def __getitem__(self, i):
308 """
309 usual
310 """
311 return self.gens[i]
313 def __iter__(self):
314 """
315 usual
316 """
317 return self.gens.__iter__()
319 def __len__(self):
320 """
321 usual
322 """
323 return len(self.gens)
325 def count(self):
326 """
327 return the distribution of healthy, sick, cured people
328 """
329 return Counter(map(lambda p: p.state, self))
331 def evolution(self, dt=0.5):
332 """
333 new iteration
335 @param dt dt
336 @return nb1,nb2
338 We walk through everybody and call
339 :meth:`evolution <ensae_teaching_cs.special.propragation_epidemics.Person.evolution>`.
340 """
342 # on renouvelle une certaine proportion de pates (renouvellement)
343 # tire au hasard
344 pop = copy.deepcopy(self)
345 for p in self.gens:
346 p.evolution(dt, pop)
347 return self.count()
350def display_person(self, screen, pygame):
351 """
352 display a person on a pygame screen
354 @param self Person
355 @param screen screen
356 @param pygame module pygame
357 """
358 c = Person.colors[self.state]
359 pygame.draw.rect(screen, c, pygame.Rect(
360 self.pos.x - 4, self.pos.y - 4, 8, 8))
363def display_population(self, screen, pygame, font, back_ground):
364 """
365 affichage
367 @param self Person
368 @param screen screen
369 @param font font (pygame)
370 @param back_ground back ground color
371 @param pygame module pygame
372 """
373 screen.fill(back_ground)
374 for p in self.gens:
375 display_person(p, screen, pygame)
376 c = self.count()
377 text = "vie: %d" % c.get(0, 0)
378 text = font.render(text, True, (255, 255, 255))
379 screen.blit(text, (self.cote, 100))
380 text = "malade: %d" % c.get(1, 0)
381 text = font.render(text, True, (255, 255, 255))
382 screen.blit(text, (self.cote, 135))
383 text = "mort: %d" % c.get(2, 0)
384 text = font.render(text, True, (255, 255, 255))
385 screen.blit(text, (self.cote, 170))
386 text = "gueris: %d" % c.get(3, 0)
387 text = font.render(text, True, (255, 255, 255))
388 screen.blit(text, (self.cote, 205))
391def pygame_simulation(pygame, first_click=False, folder=None,
392 iter=1000, cote=600, nb=(200, 20), flags=0,
393 **params):
394 """
395 Runs a graphic simulation. The user can see a :epkg:`pygame`
396 screen showing the evolution of population.
397 A healthy person is white, green is sick,
398 blue is healed, black is dead. The function can save an image for
399 every iteration. They can be merged into a video with
400 function @see fn make_video.
402 @param pygame module pygame (avoids importing in this file)
403 @param first_click starts the simulation after a first click
404 @param folder to save the simulation, an image per simulation
405 @param iter number of iterations to run
406 @param cote @see cl EpidemicPopulation
407 @param nb @see cl EpidemicPopulation
408 @param params @see cl EpidemicPopulation
409 @param flags see `pygame.display.set_mode <https://www.pygame.org/docs/ref/display.html#pygame.display.set_mode>`_
410 @param fLOG logging function
412 The simulation looks like this:
414 .. raw:: html
416 <video autoplay="" controls="" loop="" height="400">
417 <source src="http://www.xavierdupre.fr/enseignement/complements/epidemic.mp4" type="video/mp4" />
418 </video>
420 Pour lancer la simulation::
422 from ensae_teaching_cs.special.propagation_epidemic import pygame_simulation
423 import pygame
424 pygame_simulation(pygame)
426 """
427 pygame.init()
428 size = cote + 200, cote
429 screen = pygame.display.set_mode(size, flags)
430 font = pygame.font.Font("freesansbold.ttf", 30)
431 back_ground = (128, 128, 128)
433 pop = EpidemicPopulation(cote, nb)
434 display_population(pop, screen, pygame, font, back_ground)
435 pygame.display.flip()
436 if first_click:
437 wait_event(pygame)
439 for i in range(0, iter):
441 empty_main_loop(pygame)
442 nb = pop.evolution()
443 display_population(pop, screen, pygame, font, back_ground)
444 pygame.display.flip()
445 pygame.event.peek()
446 if folder is not None:
447 image = os.path.join(folder, "image_%04d.png" % i)
448 pygame.image.save(screen, image)
449 pygame.time.wait(50)
450 if 1 not in nb or nb[1] == 0:
451 break
453 if first_click:
454 wait_event(pygame)
457def numerical_simulation(nb=(200, 20), cote=600, iter=1000, fLOG=noLOG, **params):
458 """
459 Run a simulation, @see cl EpidemicPopulation.
461 @param iter number of iterations to run
462 @param cote @see cl EpidemicPopulation
463 @param nb @see cl EpidemicPopulation
464 @param params @see cl EpidemicPopulation
465 @param fLOG to display status every 10 iterations
466 @return population count
467 """
468 pop = EpidemicPopulation(cote, nb, **params)
470 lasti = None
471 for i in range(0, iter):
472 nb = pop.evolution()
473 lasti = i
474 if 1 not in nb or nb[1] == 0:
475 break
476 if i % 10 == 0:
477 fLOG("iteration", i, ":", nb)
479 r = pop.count()
480 return r, lasti