def test_generate_scenario_future_from_last_known_date_freeze(self):
# Simulate Italy freezes NPIS for the rest of the year
countries = ["Italy"]
start_date_str = None
end_date_str = "2020-12-31"
scenario = "Freeze"
last_known_date = self.latest_df[self.latest_df.CountryName ==
"Italy"].Date.max()
frozen_npis_df = self.latest_df[
(self.latest_df.CountryName == "Italy") &
(self.latest_df.Date == last_known_date)][NPI_COLUMNS].reset_index(
drop=True)
scenario_npis = list(frozen_npis_df.values[0])
# Generate the scenario
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check it
self._check_future(
start_date_str=start_date_str,
end_date_str=end_date_str,
scenario_df=scenario_df[scenario_df.CountryName == countries[0]],
scenario_npis=scenario_npis,
country=countries[0])
def test_generate_scenario_mind_the_gap_custom(self):
# Scenario = Custom
start_date_str = "2021-01-01"
end_date_str = "2021-01-31"
countries = ["Italy"]
# Set all the NPIs to one for each day between start data and end date.
scenario = [ONE_NPIS] * 31
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 2020-01-01. 366 days for 2020 + 31 for Jan 2021
self.assertEqual(
397, len(scenario_df),
"Expected the number of days between inception and end date")
# The last 31 rows must be the same
self.assertEqual(
1,
scenario_df.tail(31)[NPI_COLUMNS].mean().mean(),
"Expected the last 31 rows to have all NPIs set to 1")
def test_generate_scenario_all_countries_future_from_last_known_date_freeze(
self):
# Simulate ALL countries uses custom NPIs for the rest of the year
countries = None
end_date_str = "2020-12-31"
# Make sure we generate scenarios for enough days
nb_days = 180
scenario = [ONE_NPIS] * nb_days
# Generate the scenarios
scenario_df = generate_scenario(None,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check them
all_countries = self.latest_df.CountryName.unique()
for country in all_countries:
all_regions = self.latest_df[self.latest_df.CountryName ==
country].RegionName.unique()
for region in all_regions:
self._check_future(start_date_str=None,
end_date_str=end_date_str,
scenario_df=scenario_df[
(scenario_df.CountryName == country)
& (scenario_df.RegionName == region)],
scenario_npis=scenario[0],
country=country,
region=region)
def test_generate_scenario_future_freeze(self):
# Simulate Italy froze it's NPIS for the second part of the year
countries = ["Italy"]
start_date_str = "2020-07-01"
end_date_str = "2020-12-31"
scenario = "Freeze"
before_day = pd.to_datetime("2020-06-30", format='%Y-%m-%d')
frozen_npis_df = self.latest_df[
(self.latest_df.CountryName == "Italy")
& (self.latest_df.Date == before_day)][NPI_COLUMNS].reset_index(
drop=True)
scenario_npis = list(frozen_npis_df.values[0])
# Generate the scenario
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check it
self._check_future(
start_date_str=start_date_str,
end_date_str=end_date_str,
scenario_df=scenario_df[scenario_df.CountryName == countries[0]],
scenario_npis=scenario_npis,
country=countries[0])
def test_generate_scenario_historical(self):
latest_df = _get_dataset()
start_date_str = "2020-08-01"
end_date_str = "2020-08-4"
countries = ["Italy"]
scenario_df = generate_scenario(start_date_str, end_date_str,
latest_df, countries)
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
217, len(scenario_df),
"Expected the number of days between inception and end date")
# Check multiple countries
countries = ["France", "Italy"]
scenario_df = generate_scenario(start_date_str, end_date_str,
latest_df, countries)
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
217 * 2, len(scenario_df),
"Expected the number of days between inception and end date")
# All countries: do not pass a countries list
scenario_df = generate_scenario(start_date_str, end_date_str,
latest_df)
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(latest_df.CountryName.unique(),
scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
# Contains the regions too. -1 to remove the NaN region, already counted as a country
nb_geos = len(latest_df.CountryName.unique()) + len(
latest_df.RegionName.unique()) - 1
self.assertEqual(
217 * nb_geos, len(scenario_df),
"Expected the number of days between inception and end date")
def _check_counterfactual(self, scenario, scenario_npis):
# Simulate Italy lifted all NPI for this period
start_date_str = "2020-03-20"
end_date_str = "2020-03-26"
countries = ["Italy"]
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
self.assertFalse(scenario_df["Date"].duplicated().any(),
"Expected 1 row per date only")
start_date = pd.to_datetime(start_date_str, format='%Y-%m-%d')
end_date = pd.to_datetime(end_date_str, format='%Y-%m-%d')
before_day = start_date - np.timedelta64(1, 'D')
before_day_npis = scenario_df[
scenario_df.Date == before_day][NPI_COLUMNS].reset_index(drop=True)
before_day_npis_truth = self.latest_df[
(self.latest_df.CountryName == "Italy")
& (self.latest_df.Date == before_day)][NPI_COLUMNS].reset_index(
drop=True)
# Check the day before the scenario is correct
pd.testing.assert_frame_equal(before_day_npis_truth, before_day_npis,
"Not the expected frozen NPIs")
# For the right period (+1 to include start and end date)
nb_days = (end_date - start_date).days + 1
for i in range(nb_days):
check_day = start_date + np.timedelta64(i, 'D')
check_day_npis_df = scenario_df[
scenario_df.Date == check_day][NPI_COLUMNS].reset_index(
drop=True)
check_day_npis = list(check_day_npis_df.values[0])
self.assertListEqual(scenario_npis, check_day_npis)
# Check Mar 27 is different from frozen day
after_day = end_date + np.timedelta64(1, 'D')
after_day_npis_df = scenario_df[
scenario_df.Date == after_day][NPI_COLUMNS].reset_index(drop=True)
self.assertTrue((scenario_npis - after_day_npis_df.values[0]).any(),
"Expected NPIs to be different")
# Check 27 is indeed equal to truth
after_day_npis_truth = self.latest_df[
(self.latest_df.CountryName == "Italy")
& (self.latest_df.Date == after_day)][NPI_COLUMNS].reset_index(
drop=True)
pd.testing.assert_frame_equal(after_day_npis_truth, after_day_npis_df,
"Not the expected unfrozen NPIs")
def test_generate_scenario_mind_the_gap_freeze_2_countries(self):
# Check 2 countries
start_date_str = "2021-01-01"
end_date_str = "2021-01-31"
countries = ["France", "Italy"]
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario="Freeze")
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 2020-01-01. 366 days for 2020 + 31 for Jan 2021
self.assertEqual(
397 * 2, len(scenario_df),
"Expected the number of days between inception and end date")
def test_generate_scenario_mind_the_gap_min(self):
# Scenario = MIN
start_date_str = "2021-01-01"
end_date_str = "2021-01-31"
countries = ["Italy"]
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario="MIN")
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 2020-01-01. 366 days for 2020 + 31 for Jan 2021
self.assertEqual(
397, len(scenario_df),
"Expected the number of days between inception and end date")
# The last 31 rows must be the same
self.assertEqual(0,
scenario_df.tail(31)[NPI_COLUMNS].sum().sum(),
"Expected the last 31 rows to have NPIs set to 0")
def test_generate_scenario_future_from_last_known_date_max(self):
# Simulate Italy maxes out NPIs for the rest of the year
countries = ["Italy"]
start_date_str = None
end_date_str = "2020-12-31"
scenario = "MAX"
scenario_npis = MAX_NPIS
# Generate the scenario
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check it
self._check_future(
start_date_str=start_date_str,
end_date_str=end_date_str,
scenario_df=scenario_df[scenario_df.CountryName == countries[0]],
scenario_npis=scenario_npis,
country=countries[0])
def test_generate_scenario_future_max(self):
# Simulate Italy maxed out all NPIs for a period
countries = ["Italy"]
start_date_str = "2020-07-01"
end_date_str = "2020-12-31"
scenario = "MAX"
scenario_npis = MAX_NPIS
# Generate the scenario
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check it
self._check_future(
start_date_str=start_date_str,
end_date_str=end_date_str,
scenario_df=scenario_df[scenario_df.CountryName == countries[0]],
scenario_npis=scenario_npis,
country=countries[0])
def test_generate_scenario_future_custom(self):
# Simulate Italy used custom NPIs for a period: each NPI set to 1 for 7 consecutive days
countries = ["Italy"]
start_date_str = "2020-07-01"
end_date_str = "2020-12-31"
start_date = pd.to_datetime(start_date_str, format='%Y-%m-%d')
end_date = pd.to_datetime(end_date_str, format='%Y-%m-%d')
nb_days = (end_date - start_date).days + 1 # +1 to include start date
scenario = [ONE_NPIS] * nb_days
# Generate the scenario
scenario_df = generate_scenario(start_date_str,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check it
self._check_future(
start_date_str=start_date_str,
end_date_str=end_date_str,
scenario_df=scenario_df[scenario_df.CountryName == countries[0]],
scenario_npis=scenario[0],
country=countries[0])
def test_generate_scenario_future_from_last_known_date_custom(self):
# Simulate Italy uses custom NPIs for the rest of the year
countries = ["Italy"]
last_known_date = self.latest_df[self.latest_df.CountryName ==
"Italy"].Date.max()
start_date = last_known_date + np.timedelta64(1, 'D')
end_date_str = "2020-12-31"
end_date = pd.to_datetime(end_date_str, format='%Y-%m-%d')
nb_days = (end_date - start_date).days + 1 # +1 to include start date
scenario = [ONE_NPIS] * nb_days
# Generate the scenario
scenario_df = generate_scenario(None,
end_date_str,
self.latest_df,
countries,
scenario=scenario)
# Check it
self._check_future(start_date_str=None,
end_date_str=end_date_str,
scenario_df=scenario_df,
scenario_npis=scenario[0],
country=countries[0])
def test_generate_scenario_future(self):
# Scenario = Freeze
latest_df = _get_dataset()
start_date_str = "2021-01-01"
end_date_str = "2021-01-31"
countries = ["Italy"]
scenario_df = generate_scenario(start_date_str,
end_date_str,
latest_df,
countries,
scenario="Freeze")
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
397, len(scenario_df),
"Expected the number of days between inception and end date")
# The last 31 rows must be the same
self.assertEqual(
0,
scenario_df.tail(31)[NPI_COLUMNS].diff().sum().sum(),
"Expected the last 31 rows to have the same frozen IP")
# Scenario = MIN
scenario_df = generate_scenario(start_date_str,
end_date_str,
latest_df,
countries,
scenario="MIN")
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
397, len(scenario_df),
"Expected the number of days between inception and end date")
# The last 31 rows must be the same
self.assertEqual(0,
scenario_df.tail(31)[NPI_COLUMNS].sum().sum(),
"Expected the last 31 rows to have NPIs set to 0")
# Scenario = MAX
scenario_df = generate_scenario(start_date_str,
end_date_str,
latest_df,
countries,
scenario="MAX")
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
397, len(scenario_df),
"Expected the number of days between inception and end date")
# The last 31 rows must be the same
self.assertEqual(
sum(MAX_NPIS),
scenario_df.tail(31)[NPI_COLUMNS].mean().sum(),
"Expected the last 31 rows to have NPIs set to their max value")
# Scenario = Custom
scenario_df = generate_scenario(start_date_str,
end_date_str,
latest_df,
countries,
scenario=ONE_NPIS)
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
397, len(scenario_df),
"Expected the number of days between inception and end date")
# The last 31 rows must be the same
self.assertEqual(
1,
scenario_df.tail(31)[NPI_COLUMNS].mean().mean(),
"Expected the last 31 rows to have all NPIs set to 1")
# Check 2 countries
countries = ["France", "Italy"]
scenario_df = generate_scenario(start_date_str,
end_date_str,
latest_df,
countries,
scenario="Freeze")
self.assertIsNotNone(scenario_df)
# Misleading name but checks the elements, regardless of order
self.assertCountEqual(countries, scenario_df.CountryName.unique(),
"Not the requested countries")
# Inception is 1/1/2020
self.assertEqual(
397 * 2, len(scenario_df),
"Expected the number of days between inception and end date")