Merge pull request #110 from REAM-lab/feature/post_process

Refactoring and simplification of postprocess

Author: pesap

switch_model/wecc/get_inputs/cli.py

@@ -1,13 +1,15 @@
 """ Script to retrieve the input data from the switch-wecc database and apply post-processing steps.
 """
 import argparse
+import importlib
 import os
 
 from switch_model.utilities import query_yes_no, StepTimer
 from switch_model.wecc.get_inputs.get_inputs import query_db
-from switch_model.wecc.get_inputs.register_post_process import run_post_process
 from switch_model.wecc.utilities import load_config
-from switch_model.wecc.get_inputs.post_process_steps import *
+
+# from switch_model.wecc.get_inputs.post_process_steps import *
+# from switch_model.wecc.get_inputs.register_post_process import run_post_process, _registered_steps
 
 
 def main():
@@ -21,23 +23,64 @@ def main():
         config.yaml specifies the scenario parameters.
         The environment variable DB_URL specifies the url to connect to the database. """,
     )
-    parser.add_argument("--skip-cf", default=False, action='store_true',
-                        help="Skip creation variable_capacity_factors.csv. Useful when debugging and one doesn't"
-                             "want to wait for the command.")
-    parser.add_argument("--post-process", default=None, help="Run only this post process step.")
-    parser.add_argument("--overwrite", default=False, action='store_true',
-                        help="Overwrite previous input files without prompting to confirm.")
+    parser.add_argument(
+        "--skip-cf",
+        default=False,
+        action="store_true",
+        help="Skip creation variable_capacity_factors.csv. Useful when debugging and one doesn't"
+        "want to wait for the command.",
+    )
+    parser.add_argument(
+        "--post-process-only",
+        default=False,
+        action="store_true",
+        help="Run only post process steps.",
+    )
+    parser.add_argument(
+        "--post-process-step", default=None, help="Run only this post process step."
+    )
+    parser.add_argument(
+        "--overwrite",
+        default=False,
+        action="store_true",
+        help="Overwrite previous input files without prompting to confirm.",
+    )
     args = parser.parse_args()  # Makes switch get_inputs --help works
 
     # Load values from config.yaml
     full_config = load_config()
     switch_to_input_dir(full_config, overwrite=args.overwrite)
 
-    if args.post_process is None:
+    if not args.post_process_only and args.post_process_step is None:
         query_db(full_config, skip_cf=args.skip_cf)
-    print("Post-processing...")
-    run_post_process(full_config, step_name=args.post_process)
-    print(f"\nScript took {timer.step_time_as_str()} seconds to build input tables.")
+
+    print("\nRunning post processing...")
+
+    # Get location of post process scripts
+    post_process_path = ".".join(__name__.split(".")[:-1]) + ".post_process_steps"
+
+    def run_post_process(module):
+        """ Run a function from a given module """
+
+        # This uses python module syntax with a dot. Example: import foo.bar.test
+        mod = importlib.import_module(f".{module}", post_process_path)
+
+        post_process = getattr(mod, "post_process")
+
+        # Get specific configuration for the post process if specified
+        post_config = full_config.get(module, None)
+
+        # Run post process
+        post_process(full_config, post_config)
+
+    # Run all post process specified, otherwise run single one
+    if args.post_process_step is None:
+        for module in full_config["post_process"]:
+            run_post_process(module)
+    else:
+        run_post_process(getattr(args, "post_process_step"))
+
+    print(f"\nScript took {timer.step_time_as_str()} seconds.")
 
 
 def switch_to_input_dir(config, overwrite):
@@ -49,7 +92,7 @@ def switch_to_input_dir(config, overwrite):
         print("Inputs directory created.")
     else:
         if not overwrite and not query_yes_no(
-                "Inputs directory already exists. Allow contents to be overwritten?"
+            "Inputs directory already exists. Allow contents to be overwritten?"
         ):
             raise SystemExit("User cancelled run.")
 

switch_model/wecc/get_inputs/get_inputs.py

@@ -21,6 +21,7 @@
 from switch_model.wecc.utilities import connect
 from switch_model.version import __version__
 
+
 def write_csv_from_query(cursor, fname: str, headers: List[str], query: str):
     """Create CSV file from cursor."""
     print(f"\t{fname}.csv... ", flush=True, end="")
@@ -334,7 +335,7 @@ def query_db(config, skip_cf):
         [
             "trans_capital_cost_per_mw_km",
             "trans_lifetime_yrs",
-            "trans_fixed_om_fraction"
+            "trans_fixed_om_fraction",
         ],
         f"""
         SELECT trans_capital_cost_per_mw_km,
@@ -433,7 +434,7 @@ def query_db(config, skip_cf):
             "gen_self_discharge_rate",
             "gen_discharge_efficiency",
             "gen_land_use_rate",
-            "gen_storage_energy_to_power_ratio"
+            "gen_storage_energy_to_power_ratio",
         ],
         f"""
             select
@@ -480,7 +481,12 @@ def query_db(config, skip_cf):
     write_csv_from_query(
         db_cursor,
         "gen_build_predetermined",
-        ["GENERATION_PROJECT", "build_year", "gen_predetermined_cap", "gen_predetermined_storage_energy_mwh"],
+        [
+            "GENERATION_PROJECT",
+            "build_year",
+            "gen_predetermined_cap",
+            "gen_predetermined_storage_energy_mwh",
+        ],
         f"""select generation_plant_id, build_year, capacity as gen_predetermined_cap, gen_predetermined_storage_energy_mwh
                 from generation_plant_existing_and_planned
                 join generation_plant as t using(generation_plant_id)
@@ -529,7 +535,7 @@ def query_db(config, skip_cf):
     ########################################################
     # FINANCIALS
 
-    #updated from $2016 and 7%
+    # updated from $2016 and 7%
     write_csv(
         [[2018, 0.05, 0.05]],
         "financials",
@@ -591,7 +597,7 @@ def query_db(config, skip_cf):
         WHERE time_sample_id = {time_sample_id}
             AND study_timeframe_id = {study_timeframe_id}
         ORDER BY 1;
-        """
+        """,
     )
 
     write_csv_from_query(
@@ -655,7 +661,7 @@ def query_db(config, skip_cf):
         where period!=0
         group by period
         order by 1;
-        """
+        """,
     )
 
     ########################################################
@@ -679,7 +685,7 @@ def query_db(config, skip_cf):
             where period!=0
             group by load_zone, period
             order by 1, 2;
-            """
+            """,
         )
         modules.append("switch_model.policies.rps_unbundled")
 
@@ -901,11 +907,17 @@ def ca_policies(db_cursor, scenario_params):
     write_csv_from_query(
         db_cursor,
         "ca_policies",
-        ['PERIOD', 'ca_min_gen_timepoint_ratio', 'ca_min_gen_period_ratio', 'carbon_cap_tco2_per_yr_CA'],
-        query
+        [
+            "PERIOD",
+            "ca_min_gen_timepoint_ratio",
+            "ca_min_gen_period_ratio",
+            "carbon_cap_tco2_per_yr_CA",
+        ],
+        query,
     )
 
-    modules.append('switch_model.policies.CA_policies')
+    modules.append("switch_model.policies.CA_policies")
+
 
 def planning_reserves(db_cursor, scenario_params):
     # reserve_capacity_value.csv specifies the capacity factors that should be used when calculating
@@ -917,7 +929,7 @@ def planning_reserves(db_cursor, scenario_params):
     write_csv_from_query(
         db_cursor,
         "reserve_capacity_value",
-        ["GENERATION_PROJECT","timepoint","gen_capacity_value"],
+        ["GENERATION_PROJECT", "timepoint", "gen_capacity_value"],
         f"""
         select
             generation_plant_id,
@@ -954,7 +966,11 @@ def planning_reserves(db_cursor, scenario_params):
     write_csv_from_query(
         db_cursor,
         "planning_reserve_requirements",
-        ["PLANNING_RESERVE_REQUIREMENT", "prr_cap_reserve_margin", "prr_enforcement_timescale"],
+        [
+            "PLANNING_RESERVE_REQUIREMENT",
+            "prr_cap_reserve_margin",
+            "prr_enforcement_timescale",
+        ],
         """
         SELECT
             planning_reserve_requirement, prr_cap_reserve_margin, prr_enforcement_timescale

switch_model/wecc/get_inputs/post_process_steps/add_storage.py

@@ -5,7 +5,6 @@
 the csvs in the inputs folder.
 """
 import pandas as pd
-
 from switch_model.wecc.get_inputs.register_post_process import register_post_process
 
 
@@ -17,16 +16,16 @@ def fetch_df(tab_name, key, config):
         "constants": 0,
         "plants": 889129113,
         "costs": 1401952285,
-        "minimums": 1049456965
+        "minimums": 1049456965,
     }
     SHEET_ID = "1SJrj039T1T95NLTs964VQnsfZgo2QWCo29x2ireVYcU"
 
     gid = TAB_NAME_GID[tab_name]
     url = f"https://docs.google.com/spreadsheet/ccc?key={SHEET_ID}&output=csv&gid={gid}"
 
-    df: pd.DataFrame = pd.read_csv(url, index_col=False) \
-        .replace("FALSE", False) \
-        .replace("TRUE", True)
+    df: pd.DataFrame = (
+        pd.read_csv(url, index_col=False).replace("FALSE", False).replace("TRUE", True)
+    )
 
     if "description" in df.columns:
         df = df.drop("description", axis=1)
@@ -43,17 +42,16 @@ def filer_by_scenario(df, scenario_column, config):
     if scenario_column in config:
         scenario = config[scenario_column]
     else:
-        scenario = input(f"Which scenario do you want for '{scenario_column}' (default 0) : ")
+        scenario = input(
+            f"Which scenario do you want for '{scenario_column}' (default 0) : "
+        )
         scenario = int(scenario) if scenario != "" else 0
     df = df[df[scenario_column] == scenario]
     return df.drop(scenario_column, axis=1)
 
 
 def cross_join(df1, df2):
-    return df1.assign(key=1).merge(
-        df2.assign(key=1),
-        on="key"
-    ).drop("key", axis=1)
+    return df1.assign(key=1).merge(df2.assign(key=1), on="key").drop("key", axis=1)
 
 
 def add_to_csv(filename, to_add, primary_key=None, append=True):
@@ -83,8 +81,12 @@ def drop_previous_candidate_storage():
 
     gen = pd.read_csv("generation_projects_info.csv", index_col=False)
     # Find generation projects that are both storage and not predetermined (i.e. candidate)
-    predetermined_gen = pd.read_csv("gen_build_predetermined.csv", index_col=False)["GENERATION_PROJECT"]
-    should_drop = (gen["gen_tech"] == STORAGE_TECH) & ~gen["GENERATION_PROJECT"].isin(predetermined_gen)
+    predetermined_gen = pd.read_csv("gen_build_predetermined.csv", index_col=False)[
+        "GENERATION_PROJECT"
+    ]
+    should_drop = (gen["gen_tech"] == STORAGE_TECH) & ~gen["GENERATION_PROJECT"].isin(
+        predetermined_gen
+    )
     # Find projects that we should drop (candidate storage)
     gen_to_drop = gen[should_drop]["GENERATION_PROJECT"]
 
@@ -99,30 +101,46 @@ def drop_previous_candidate_storage():
 
 
 @register_post_process(
-    name="add_storage",
     msg="Adding storage from Google Sheets",
-    only_with_config=True,
-    priority=1  # Increased priority (default is 2) so that it always runs before replace_plants_in_zone_all.py
 )
-def main(config):
+def post_process(config):
     # Drop previous candidate storage from inputs
     drop_previous_candidate_storage()
 
     # Get the generation storage plants from Google Sheet
-    gen_projects = fetch_df("constants", "constant_scenario", config).set_index("param_name").transpose()
-    gen_projects = cross_join(gen_projects, fetch_df("plants", "plants_scenario", config))
+    gen_projects = (
+        fetch_df("constants", "constant_scenario", config)
+        .set_index("param_name")
+        .transpose()
+    )
+    gen_projects = cross_join(
+        gen_projects, fetch_df("plants", "plants_scenario", config)
+    )
 
     # Append the storage plants to the inputs
-    add_to_csv("generation_projects_info.csv", gen_projects, primary_key="GENERATION_PROJECT")
+    add_to_csv(
+        "generation_projects_info.csv", gen_projects, primary_key="GENERATION_PROJECT"
+    )
 
     # Create min_per_tech.csv
     min_projects = fetch_df("minimums", "minimums_scenario", config)
-    add_to_csv("min_per_tech.csv", min_projects, primary_key=["gen_tech", "period"], append=False)
+    add_to_csv(
+        "min_per_tech.csv",
+        min_projects,
+        primary_key=["gen_tech", "period"],
+        append=False,
+    )
 
     # Get the plant costs from GSheets and append to costs
     storage_costs = fetch_df("costs", "costs_scenario", config)
-    storage_costs = storage_costs[storage_costs["GENERATION_PROJECT"].isin(gen_projects["GENERATION_PROJECT"])]
-    add_to_csv("gen_build_costs.csv", storage_costs, primary_key=["GENERATION_PROJECT", "build_year"])
+    storage_costs = storage_costs[
+        storage_costs["GENERATION_PROJECT"].isin(gen_projects["GENERATION_PROJECT"])
+    ]
+    add_to_csv(
+        "gen_build_costs.csv",
+        storage_costs,
+        primary_key=["GENERATION_PROJECT", "build_year"],
+    )
 
     # Create add_storage_info.csv
     pd.DataFrame([config]).transpose().to_csv("add_storage_info.csv", header=False)
@@ -132,9 +150,9 @@ def main(config):
     gen_type.columns = ["gen_tech", "energy_source"]
     gen_type["map_name"] = "default"
     gen_type["gen_type"] = "Storage"
-    pd.concat([
-        pd.read_csv("graph_tech_types.csv", index_col=False), gen_type
-    ]).to_csv("graph_tech_types.csv", index=False)
+    pd.concat([pd.read_csv("graph_tech_types.csv", index_col=False), gen_type]).to_csv(
+        "graph_tech_types.csv", index=False
+    )
 
 
 if __name__ == "__main__":