From 5d0e6053b9326dd50bffb5ca45255d3b33aa5e03 Mon Sep 17 00:00:00 2001 From: Ben Date: Wed, 21 Apr 2021 20:39:04 -0400 Subject: [PATCH] Implemented state pattern to aid with determining if a country is given or not --- src/covidify/data_prep.py | 30 +++---- src/covidify/data_visualization.py | 139 +++++++++++++++++++++-------- 2 files changed, 118 insertions(+), 51 deletions(-) diff --git a/src/covidify/data_prep.py b/src/covidify/data_prep.py index 6153d44..e26a8b5 100644 --- a/src/covidify/data_prep.py +++ b/src/covidify/data_prep.py @@ -19,7 +19,7 @@ import pandas as pd from string import capwords from difflib import get_close_matches -from datetime import datetime, date, time +from datetime import datetime, date, time from covidify.sources import github, wiki from covidify.config import REPO, TMP_FOLDER, TMP_GIT, DATA @@ -43,21 +43,21 @@ if source == 'JHU': df = github.get() - + elif source == 'wiki': print('Apologies, the wikipedia source is not ready yet - getting github data') df = github.get() - + ############ COUNTRY SELECTION ############ def get_similar_countries(c, country_list): pos_countries = get_close_matches(c, country_list) - + if len(pos_countries) > 0: print('\033[1;31m'+c, 'was not listed. did you mean', pos_countries[0].capitalize() + '?\033[0;0m') - + #Only delete if its a covidify generated folder if 'Desktop/covidify-output-' in out: os.system('rm -rf ' + out) @@ -67,13 +67,13 @@ def get_similar_countries(c, country_list): if 'Desktop/covidify-output-' in out: os.system('rm -rf ' + out) sys.exit(1) - + def check_specified_country(df, country): ''' let user filter reports by country, if not found then give a option if the string is similar ''' - + # Get all unique countries in the data country_list = list(map(lambda x:x.lower().strip(), set(df.country.values))) @@ -85,7 +85,7 @@ def check_specified_country(df, country): # give similar option if similarity found if country.lower() not in country_list: get_similar_countries(country, country_list) - + else: #Return filtered dataframe print('... filtering data for', country) @@ -170,9 +170,9 @@ def get_top_countries(data): # Get top N infected countries tmp_df = data.copy() tmp_df = tmp_df[tmp_df.file_date == df.file_date.max()] - return tmp_df.groupby(['country']).agg({'confirmed': 'sum'}).sort_values('confirmed',ascending=False).head(top).index - -TOP_N_COUNTRIES = get_top_countries(df) + return tmp_df.groupby(['country']).agg({'confirmed': 'sum'}).sort_values('confirmed',ascending=False).head(top).index + +TOP_N_COUNTRIES = get_top_countries(df) tmp_df = df[df.country.isin(TOP_N_COUNTRIES)].copy() @@ -188,7 +188,7 @@ def get_day_counts(d, country): 'deaths': 'sum'}) result_df['date'] = data['file_date'].unique() result_df['country'] = country - + result_df = result_df[result_df.confirmed >= 500] result_df.insert(loc=0, column='day', value=np.arange(len(result_df))) return result_df @@ -196,10 +196,10 @@ def get_day_counts(d, country): df_list = [] for country in TOP_N_COUNTRIES: - print(' ...', country + ': ' + str(tmp_df[(tmp_df.file_date == df.file_date.max()) & + print(' ...', country + ': ' + str(tmp_df[(tmp_df.file_date == df.file_date.max()) & (tmp_df.country == country)].confirmed.sum())) df_list.append(get_day_counts(tmp_df[tmp_df.country == country], country)) - + log_df = pd.concat(df_list, axis=0, ignore_index=True) @@ -227,4 +227,4 @@ def get_day_counts(d, country): log_df.astype(str).to_csv(os.path.join(save_dir, log_file_name)) print('...', log_file_name) -print('Done!') \ No newline at end of file +print('Done!') diff --git a/src/covidify/data_visualization.py b/src/covidify/data_visualization.py index 8674540..9b45592 100644 --- a/src/covidify/data_visualization.py +++ b/src/covidify/data_visualization.py @@ -22,13 +22,14 @@ from datetime import datetime import matplotlib.pyplot as plt from covidify.utils.utils import replace_arg_score +import abc # plt settings font = {'weight' : 'bold', 'size' : 22} plt.rc('font', **font) plt.style.use('ggplot') - + args = docopt.docopt(__doc__) out = args['--output_folder'] country = args['--country'] @@ -41,19 +42,15 @@ if country == 'Global': country = None -#change report name if country specified -def create_report_name(country): - if country: - return '{}_report_{}.xlsx'.format(country, datetime.date(datetime.now())) - else: - return 'report_{}.xlsx'.format(datetime.date(datetime.now())) +#Create instance of Country class +context_country = Country(country) # Dynamic parameters data_dir = os.path.join(out, 'data', str(datetime.date(datetime.now()))) agg_file = 'agg_data_{}.csv'.format(datetime.date(datetime.now())) trend_file = 'trend_{}.csv'.format(datetime.date(datetime.now())) log_file = 'log_{}.csv'.format(datetime.date(datetime.now())) -report = create_report_name(country) +report = context_country.create_report_name(country) # import data @@ -75,54 +72,123 @@ def create_report_name(country): agg_df[col] = agg_df[col].replace('', 0).astype(int) -##### Define Graphs ##### -#Change titles and saved file names if country -#is specified -def create_title(fig_title, country): - if country: + +class State(metaclass=abc.ABCMeta): +''' +State interface +''' + def __init__(self): + pass + + #Creates title used for graphs + @abc.abstractmethod + def create_title(self): + pass + + #Creates safe file name for graphs + @abc.abstractmethod + def create_save_file(self): + pass + + #Creates report name + @abc.abstractmethod + def create_report_name(self): + pass + + + +class CountryGivenState(State): +''' +State to be used if user gives a country +''' + def __init__(self): + pass + + def create_title(self, fig_title, country): return fig_title + ' for ' + country - else: - return fig_title - -def create_save_file(col, country, graph_type): - if country: + + def create_save_file(self, col, country, graph_type): return '{}_{}_{}.png'.format(country, col, graph_type) - else: - return '{}_{}.png'.format(col, graph_type) + + def create_report_name(self, country): + return '{}_report_{}.xlsx'.format(country, datetime.date(datetime.now())) + +class GlobalCountryState(State): +''' +State to be used if no country was specifed and in that case we include all countries +''' + + def __init__(self): + pass + + def create_title(self, fig_title, country): + return fig_title + + def create_save_file(self, col, country, graph_type): + return '{}_{}_{}.png'.format(col, graph_type) + + def create_report_name(self, country): + return 'report_{}.xlsx'.format(datetime.date(datetime.now())) + +class Country: +''' +Context class for the different states +This class should be the one that is instantied and called upon +''' + def __init__(self, country): + if country: + self._state = CountryGivenState() + else + self._state = GlobalCountryState() + + def create_title(self, fig_title, country): + self._state.create_title(self, fig_title, country) + + def create_save_file(self, col, country, graph_type): + self._state.create_save_file(self, col, country, graph_type) + + def create_report_name(self, country): + self._state.create_report_name(self, country) + + def change_state(self, country): + if country: + self._state = CountryGivenState() + else + self._state = GlobalCountryState() # Plot and save trendline graph def create_trend_line(tmp_df, date_col, col, col2, col3, fig_title, country): fig, ax = plt.subplots(figsize=(20,10)) tmp_df.groupby([date_col])[[col, col2, col3]].sum().plot(ax=ax, marker='o') - ax.set_title(create_title(fig_title, country)) + ax.set_title(context_country.create_title(fig_title, country)) fig = ax.get_figure() - fig.savefig(os.path.join(image_dir, create_save_file(col, country, 'trendline'))) + fig.savefig(os.path.join(image_dir, context_country.create_save_file(col, country, 'trendline'))) def create_bar(tmp_df, col, rgb, country): tmp_df = tmp_df.tail(120) fig, ax = plt.subplots(figsize=(20,10)) tmp = tmp_df.groupby(['date'])[[col]].sum() - ax.set_title(create_title(col, country)) + ax.set_title(context_country.create_title(col, country)) tmp.plot.bar(ax=ax, rot=90, color=rgb) fig = ax.get_figure() - fig.savefig(os.path.join(image_dir, create_save_file(col, country, 'bar'))) + fig.savefig(os.path.join(image_dir, context_country.create_save_file(col, country, 'bar'))) def create_stacked_bar(tmp_df, col1, col2, fig_title, country): tmp_df = tmp_df.tail(120) tmp_df = tmp_df.set_index('date') fig, ax = plt.subplots(figsize=(20,10)) - ax.set_title(create_title(fig_title, country)) + ax.set_title(context_country.create_title(fig_title, country)) tmp_df[[col2, col1]].plot.bar(ax=ax, rot=90, stacked=True) fig = ax.get_figure() - fig.savefig(os.path.join(image_dir, create_save_file(col2, country, 'stacked_bar'))) + fig.savefig(os.path.join(image_dir, context_country.create_save_file(col2, country, 'stacked_bar'))) def log_plot(tmp, col, fig_title): ''' Plot on a logarithmic scale for comparing countries infection rates - + ''' cm = plt.get_cmap('gist_rainbow') fig = plt.figure(figsize = (20,10)) @@ -136,8 +202,9 @@ def log_plot(tmp, col, fig_title): ax.set_yscale('log', basey=10) ax.set_title(fig_title) fig = ax.get_figure() - fig.savefig(os.path.join(image_dir, create_save_file(col, country=None, graph_type='log'))) - + context_country.change_state(None) + fig.savefig(os.path.join(image_dir, context_country.create_save_file(col, country=, graph_type='log'))) + ##### Create Graphs ##### print('Creating graphs...') print('... Time Series Trend Line') @@ -171,7 +238,7 @@ def log_plot(tmp, col, fig_title): workbook_writer = pd.ExcelWriter(os.path.join(reports_dir, report), engine='xlsxwriter') # Add daily summary to spreadsheet -daily_df.to_excel(workbook_writer, sheet_name='daily figures') +daily_df.to_excel(workbook_writer, sheet_name='daily figures') workbook = workbook_writer.book def get_image_types(path): @@ -182,14 +249,14 @@ def get_image_types(path): types = [] for fn in glob.glob(os.path.join(path, '*.png')): types.append(fn.split('_',)[-1].split('.')[0]) - + return types # Get all images for each type def read_images(path, graph_type): image_list = [] for fn in glob.glob(os.path.join(path, '*_{}.png'.format(graph_type))): - image_list.append(fn) + image_list.append(fn) images = {graph_type : image_list} return dict(images) @@ -199,13 +266,13 @@ def read_images(path, graph_type): for types in set(image_types): print('... reading images for:', types) type_dict = read_images(image_dir, types) - + # Add image to the worksheet worksheet = workbook.add_worksheet(name='{}_graphs'.format(types)) for image in type_dict[types]: - worksheet.insert_image('A' +str(padding), image) + worksheet.insert_image('A' +str(padding), image) padding += 50 padding = 1 - + workbook.close() -print('Done!') \ No newline at end of file +print('Done!')