Предсказание числа заболевших коронавирусом из Китая

corona.csv

coronavir_predict.py

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit
from datetime import timedelta
from prettytable import PrettyTable


# загрузка данных
df = pd.read_csv('corona.csv', parse_dates=['day'])
print(df.head())


# экспоненциальная функция с параметрами
def fit_func_exp(x, a, b, c):
    return a * np.exp(c * x) + b


# зависимая переменная y - ее будем предсказывать
train_infected = df['infected']
train_dead = df['dead']

# дни - преобразуем их в целые числа от 0 до максимального
days = range(len(df['day']))

# у нас 3 параметра в функции a, b, c – начальное приближение - единицы
p0 = np.ones(3)

# впишем кривую в наши данные
(a, b, c), _ = curve_fit(fit_func_exp, days, train_infected, p0=p0)
(a1, b1, c1), _ = curve_fit(fit_func_exp, days, train_dead, p0=p0)


# предскажем динамику вируса на 60 дней (начало соотвествует началу из данных 20 янв 2020)
MAX_DAY = 60

x_days = np.linspace(0, MAX_DAY - 1, MAX_DAY)
y_infect = fit_func_exp(x_days, a, b, c)
y_dead = fit_func_exp(x_days, a1, b1, c1)

# население Земли
EARTH_POPULATION = 7_530_000_000

# найдем номер дня, когда количество зараженных достигнет популяции Земли
doom_inf_index = np.argmax(y_infect >= EARTH_POPULATION)
doom_inf_day = x_days[doom_inf_index]

doom_dead_index = np.argmax(y_dead >= EARTH_POPULATION)
doom_dead_day = x_days[doom_dead_index]

# вычислим дату
day0 = df['day'][0]
doom_inf_date = day0 + timedelta(days=int(doom_inf_index))
doom_dead_date = day0 + timedelta(days=int(doom_dead_index))

# дата конца!
print(f'Дата заражение Земли: {doom_inf_date:%d.%m.%Y}')
print(f'Судный день: {doom_dead_date:%d.%m.%Y}')

# --- таблицу рисуем ---

table = PrettyTable()
table.field_names = ['Дата', 'Число заболевших', 'Число погибших']

cur_day = day0
for day, (inf_ppl, dead_ppl) in enumerate(zip(y_infect, y_dead)):

    table.add_row([
        f'{cur_day:%d.%m.%Y}',
        int(inf_ppl),
        int(dead_ppl)
    ])
    cur_day += timedelta(days=1)

print(table)

# ----------------------

plt.xlabel('Дни')
plt.ylabel('Больные')

# в log шкале
plt.yscale('log')

# это данные текущей статистики - нарисуем их синими точками
plt.scatter(days, train_infected, marker='D', label='Реальные')

# это красная линия – предсказание
plt.plot(x_days[:22], y_infect[:22], 'r', label='Предсказание')
plt.legend()

plt.show()


# df.plot(kind='bar', x='day', y=['infected', 'dead'])
# plt.show()

Результаты:

Doom date: 13.03.2020