read 311 pickle when necessary

PSLmodels · Oct 14, 2023 · 1bec7ed · 1bec7ed
1 parent 3fa9d50
commit 1bec7ed
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 18 deletions.
diff --git a/tests/test_output_plots.py b/tests/test_output_plots.py
@@ -17,18 +17,28 @@
 base_tpi = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TPI_vars_baseline.pkl")
 )
-base_params = utils.safe_read_pickle(
-    os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
-)
+if(os.python.version<3.11):
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
+    )
+else:
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline_v311.pkl")
+    )
 reform_ss = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "SS_vars_reform.pkl")
 )
 reform_tpi = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TPI_vars_reform.pkl")
 )
-reform_params = utils.safe_read_pickle(
-    os.path.join(CUR_PATH, "test_io_data", "model_params_reform.pkl")
-)
+if(os.python.version<3.11):
+    reform_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_reform.pkl")
+    )
+else:
+    reform_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_reform_v311.pkl")
+    )
 reform_taxfunctions = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TxFuncEst_reform.pkl")
 )

diff --git a/tests/test_output_tables.py b/tests/test_output_tables.py
@@ -17,18 +17,28 @@
 base_tpi = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TPI_vars_baseline.pkl")
 )
-base_params = utils.safe_read_pickle(
-    os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
-)
+if(os.python.version<3.11):
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
+    )
+else:
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline_v311.pkl")
+    )
 reform_ss = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "SS_vars_reform.pkl")
 )
 reform_tpi = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TPI_vars_reform.pkl")
 )
-reform_params = utils.safe_read_pickle(
-    os.path.join(CUR_PATH, "test_io_data", "model_params_reform.pkl")
-)
+if(os.python.version<3.11):
+    reform_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_reform.pkl")
+    )
+else:
+    reform_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_reform_v311.pkl")
+    )
 # add investment tax credit parameter that not in cached parameters
 base_params.inv_tax_credit = np.zeros(
     (base_params.T + base_params.S, base_params.M)

diff --git a/tests/test_parameter_plots.py b/tests/test_parameter_plots.py
@@ -13,9 +13,14 @@
 
 # Load in test results and parameters
 CUR_PATH = os.path.abspath(os.path.dirname(__file__))
-base_params = utils.safe_read_pickle(
-    os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
-)
+if(os.python.version<3.11):
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
+    )
+else:
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline_v311.pkl")
+    )
 base_taxfunctions = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TxFuncEst_baseline.pkl")
 )

diff --git a/tests/test_parameter_tables.py b/tests/test_parameter_tables.py
@@ -11,9 +11,14 @@
 # Load in test results and parameters
 CUR_PATH = os.path.abspath(os.path.dirname(__file__))
 
-base_params = utils.safe_read_pickle(
-    os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
-)
+if(os.python.version<3.11):
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline.pkl")
+    )
+else:
+    base_params = utils.safe_read_pickle(
+        os.path.join(CUR_PATH, "test_io_data", "model_params_baseline_v311.pkl")
+    )
 base_taxfunctions = utils.safe_read_pickle(
     os.path.join(CUR_PATH, "test_io_data", "TxFuncEst_baseline.pkl")
 )