Add function to apply custom growthfactors

Author: nikhilwoodruff

Makefile

@@ -14,6 +14,7 @@ format:
 test:
 	policyengine-core test policyengine_uk/tests/policy -c policyengine_uk
 	pytest policyengine_uk/tests/ -v
+	python policyengine_uk/data/economic_assumptions.py
 
 documentation:
 	jb clean docs/book

policyengine_uk/data/dataset_schema.py

@@ -28,8 +28,9 @@ def __init__(
             self.benunit = benunit
             self.household = household
 
-        self.data_format = "time_period_arrays"
+        self.data_format = "arrays"
         self.time_period = fiscal_year
+        self.tables = (self.person, self.benunit, self.household)
 
     def save(self, file_path: str):
         with pd.HDFStore(file_path) as f:
@@ -43,9 +44,7 @@ def load(self):
         data = {}
         for df in (self.person, self.benunit, self.household):
             for col in df.columns:
-                data[col] = {
-                    self.time_period: df[col].values,
-                }
+                data[col] = df[col].values
 
         return data
 

policyengine_uk/data/economic_assumptions.py

@@ -1,5 +1,6 @@
 import pandas as pd
 from pathlib import Path
+from .dataset_schema import UKDataset
 
 START_YEAR = 2020
 END_YEAR = 2034
@@ -39,8 +40,26 @@ def create_policyengine_uprating_factors_table():
         df_growth[year] = round(df_growth[year] / df_growth[year - 1] - 1, 3)
     df_growth[START_YEAR] = 0
 
-    df_growth.to_csv(Path(__file__).parent / "uprating_growth_factors.csv")
-    return df
+    file_path = Path(__file__).parent / "uprating_growth_factors.csv"
+    df_growth.to_csv(file_path)
+    return pd.read_csv(file_path)
+
+def apply_growth_factors(dataset: UKDataset, growth_factors: pd.DataFrame, start_year: int, end_year: int):
+    dataset = dataset.copy()
+    for i in range(len(growth_factors)):
+        index = 1
+        variable = growth_factors["Variable"].values[i]
+        for year in range(start_year, end_year + 1):
+            growth_factor = growth_factors[str(year)].values[i]
+            index *= 1 + growth_factor
+        
+        for table in dataset.tables:
+            if variable in table.columns:
+                table[variable] *= index
+        
+    return dataset
+
+BASELINE_GROWFACTORS = create_policyengine_uprating_factors_table()
 
 
 if __name__ == "__main__":