covidsimulator/data.py at master · ssannkkallpp/covidsimulator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import os
import json
import logging
import numpy as np
import pandas as pd

from datetime import date
from prophet import Prophet

# Timeframe for the data read is initiated from Jan 22, 2021 : 22/01/2020
_COVID_CASES_OVER_TIME_DATA = "https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_confirmed_global.csv"
_COVID_DEATHS_OVER_TIME_DATA = "https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_deaths_global.csv"
_COVID_RECOVERED_OVER_TIME_DATA = "https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_recovered_global.csv"


_DATA_PATH = "data"

_FTypes = [
    "cumulative_total",
    "cumulative_death",
    "cumulative_recovered",
    "incident_total",
    "incident_death",
    "incident_recovered",
]


class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)


def get_incident_cases(array):
    n = len(array)
    if n == 0:
        raise ValueError(
            "Empty list passed to get incident cases from, check fetch_data()\n"
        )
    res = [array[0]]
    for i in range(1, len(array)):
        res.append(array[i] - array[i - 1])

    return res


def write_to_file(name, content):
    with open(_DATA_PATH + "/" + name, "w") as f:
        for stuff in content:
            f.write("%s\n" % stuff)


def fetch_data():
    if (
        _COVID_CASES_OVER_TIME_DATA
        or _COVID_DEATHS_OVER_TIME_DATA
        or _COVID_RECOVERED_OVER_TIME_DATA
    ) is None:
        raise AssertionError(
            "Data links to online csv files not correctly configured\n"
        )

    #   print("Loading data from online repo...\n")
    logger = logging.getLogger("fbprophet")
    logger.setLevel(logging.ERROR)
    logger.propagate = False
    confirmed_df = pd.read_csv(_COVID_CASES_OVER_TIME_DATA)
    recovered_df = pd.read_csv(_COVID_RECOVERED_OVER_TIME_DATA)
    death_df = pd.read_csv(_COVID_DEATHS_OVER_TIME_DATA)

    countries = confirmed_df["Country/Region"].unique()

    confirmed_df = confirmed_df.drop(columns=["Province/State", "Lat", "Long"])
    confirmed_df = confirmed_df.groupby("Country/Region").agg("sum")

    recovered_df = recovered_df.drop(columns=["Province/State", "Lat", "Long"])
    recovered_df = recovered_df.groupby("Country/Region").agg("sum")

    death_df = death_df.drop(columns=["Province/State", "Lat", "Long"])
    death_df = death_df.groupby("Country/Region").agg("sum")

    #  print("Extracted data, now saving to dataset...\n")
    dates = list(confirmed_df.columns.values)

    if not os.path.isdir(_DATA_PATH):
        os.mkdir(_DATA_PATH)

    for c in countries:
        dir_string = _DATA_PATH + "/" + str(c)
        if not os.path.isdir(dir_string):
            os.mkdir(dir_string)

    num_countries = len(countries)
    for c in range(num_countries):
        country_name = confirmed_df.iloc[c].name
        data_table = []
        data_table.append(list(confirmed_df.iloc[c]))
        data_table.append(list(death_df.iloc[c]))
        data_table.append(list(recovered_df.iloc[c]))
        data_table.append(get_incident_cases(data_table[0]))
        data_table.append(get_incident_cases(data_table[1]))
        data_table.append(get_incident_cases(data_table[2]))

        prophet_df = pd.DataFrame(dates)
        prophet_df.columns = ["ds"]
        date_today = prophet_df.iloc[-1, prophet_df.columns.get_loc("ds")]
        for idx in range(len(_FTypes)):
            file = open(_DATA_PATH + "/" + country_name + "/" + _FTypes[idx], "w")
            np.savetxt(file, data_table[idx])

            prophet_df["y"] = data_table[idx]
            model = Prophet()
            model.fit(prophet_df)
            future = model.make_future_dataframe(periods=60)
            forecast = model.predict(future)
            forecast = forecast[~(forecast["ds"] < date_today)]
            forecast = forecast.iloc[1:]
            file.close()
            predicted_file = open(
                _DATA_PATH + "/" + country_name + "/predicted_" + _FTypes[idx], "w"
            )
            np.savetxt(predicted_file, forecast["yhat"].tolist())
            predicted_file.close()
            if c == 0 and idx == 0:
                predicted_dates = []
                for timeobject in forecast["ds"].tolist():
                    predicted_dates.append(timeobject.strftime("%m/%d/%y"))

                write_to_file("predicted_dates", predicted_dates)

    write_to_file("dates", dates)
    write_to_file("countries", countries)
    #  print("Done loading and saving data set\n")
    return []


if __name__ == "__main__":
    fetch_data()