diff --git a/covsirphy/analysis/data_handler.py b/covsirphy/analysis/data_handler.py index 20cbcde86..8a6d9d7f3 100644 --- a/covsirphy/analysis/data_handler.py +++ b/covsirphy/analysis/data_handler.py @@ -49,6 +49,10 @@ def __init__(self, country, province=None, **kwargs): self._area_dict = {"country": str(country), "province": str(province or self.UNKNOWN)} # Data {str: instance} self._data_dict = dict.fromkeys(self.MAIN_DICT.keys(), None) + # Raw data without complement + self._subset_df_raw = None + # With complement + self._subset_df = None # Population self._population = None # Auto complement: manually changed with DataHandler.switch_complement() @@ -218,15 +222,16 @@ def records_main(self): if jhu_data is None: raise NotRegisteredMainError(".register(jhu_data)") # Subsetting - df, self._complemented = jhu_data.records( - **self._area_dict, - start_date=self._first_date, end_date=self._last_date, - population=self._population, - **self._complement_dict, - ) + if self._subset_df is None: + self._subset_df, self._complemented = jhu_data.records( + **self._area_dict, + start_date=self._first_date, end_date=self._last_date, + population=self._population, + **self._complement_dict, + ) # Columns which are included in the main dataset except for 'Date' - self._main_cols = list(set(df.columns) - set([self.DATE])) - return df + self._main_cols = list(set(self._subset_df.columns) - set([self.DATE])) + return self._subset_df def switch_complement(self, whether=None, **kwargs): """