Minor refactor to allow running with 'python -m ...'

staadecker · staadecker · commit 53dce78e2249 · 2023-01-27T23:40:33.000-06:00
diff --git a/switch_model/__main__.py b/switch_model/__main__.py
@@ -37,7 +37,7 @@ def runner():
     "solve-scenarios": get_module_runner("switch_model.solve_scenarios"),
     "test": get_module_runner("switch_model.test"),
     "upgrade": get_module_runner("switch_model.upgrade"),
-    "get_inputs": get_module_runner("switch_model.wecc.get_inputs"),
+    "get_inputs": get_module_runner("switch_model.wecc.get_inputs.cli"),
     "drop": get_module_runner("switch_model.tools.drop"),
     "new": get_module_runner("switch_model.tools.new"),
     "graph": get_module_runner("switch_model.tools.graph.cli_graph"),
diff --git a/switch_model/wecc/get_inputs/__init__.py b/switch_model/wecc/get_inputs/__init__.py
@@ -1,71 +0,0 @@
-""" Script to retrieve the input data from the switch-wecc database and apply post-processing steps.
-"""
-import argparse
-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 *
-
-
-def main():
-    timer = StepTimer()
-
-    # Create command line tool, just provides help information
-    parser = argparse.ArgumentParser(
-        description="Write SWITCH input files from database tables.",
-        epilog="""
-        This tool will populate the inputs folder with the data from the PostgreSQL database.
-        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.",
-    )
-    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:
-        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.")
-
-
-def switch_to_input_dir(config, overwrite):
-    inputs_dir = config["inputs_dir"]
-
-    # Create inputs_dir if it doesn't exist
-    if not os.path.exists(inputs_dir):
-        os.makedirs(inputs_dir)
-        print("Inputs directory created.")
-    else:
-        if not overwrite and not query_yes_no(
-            "Inputs directory already exists. Allow contents to be overwritten?"
-        ):
-            raise SystemExit("User cancelled run.")
-
-    os.chdir(inputs_dir)
-    return inputs_dir
-
-
-if __name__ == "__main__":
-    main()
diff --git a/switch_model/wecc/get_inputs/cli.py b/switch_model/wecc/get_inputs/cli.py
@@ -0,0 +1,71 @@
+""" Script to retrieve the input data from the switch-wecc database and apply post-processing steps.
+"""
+import argparse
+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 *
+
+
+def main():
+    timer = StepTimer()
+
+    # Create command line tool, just provides help information
+    parser = argparse.ArgumentParser(
+        description="Write SWITCH input files from database tables.",
+        epilog="""
+        This tool will populate the inputs folder with the data from the PostgreSQL database.
+        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.",
+    )
+    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:
+        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.")
+
+
+def switch_to_input_dir(config, overwrite):
+    inputs_dir = config["inputs_dir"]
+
+    # Create inputs_dir if it doesn't exist
+    if not os.path.exists(inputs_dir):
+        os.makedirs(inputs_dir)
+        print("Inputs directory created.")
+    else:
+        if not overwrite and not query_yes_no(
+            "Inputs directory already exists. Allow contents to be overwritten?"
+        ):
+            raise SystemExit("User cancelled run.")
+
+    os.chdir(inputs_dir)
+    return inputs_dir
+
+
+if __name__ == "__main__":
+    main()
diff --git a/switch_model/wecc/get_inputs/post_process_steps/aggregate_candidate_projects.py b/switch_model/wecc/get_inputs/post_process_steps/aggregate_candidate_projects.py
@@ -36,7 +36,8 @@ def post_process(config):
     assert "Hydro_Pumped" not in agg_techs
 
     print(
-        f"\t\tAggregating on projects where gen_tech in {agg_techs} with capacity factors from the {cf_quantile * 100}th percentile")
+        f"\t\tAggregating on projects where gen_tech in {agg_techs} with capacity factors from the {cf_quantile * 100}th percentile"
+    )
     key = "GENERATION_PROJECT"
 
     #################
@@ -177,25 +178,39 @@ def create_capacity_factors():
         outputs/dispatch.csv (to know the DispatchGen and Curtailment)
     """
     # Read the projects
-    projects = pd.read_csv("inputs/generation_projects_info.csv",
-                           usecols=["GENERATION_PROJECT", "gen_tech", "gen_is_variable", "gen_is_baseload",
-                                    "gen_forced_outage_rate"],
-                           dtype={"GENERATION_PROJECT": str},
-                           index_col=False)
+    projects = pd.read_csv(
+        "inputs/generation_projects_info.csv",
+        usecols=[
+            "GENERATION_PROJECT",
+            "gen_tech",
+            "gen_is_variable",
+            "gen_is_baseload",
+            "gen_forced_outage_rate",
+        ],
+        dtype={"GENERATION_PROJECT": str},
+        index_col=False,
+    )
     # Filter out predetermined plants
-    predetermined = pd.read_csv("inputs/gen_build_predetermined.csv", usecols=["GENERATION_PROJECT"],
-                                dtype={"GENERATION_PROJECT": str},
-                                index_col=False)["GENERATION_PROJECT"]
+    predetermined = pd.read_csv(
+        "inputs/gen_build_predetermined.csv",
+        usecols=["GENERATION_PROJECT"],
+        dtype={"GENERATION_PROJECT": str},
+        index_col=False,
+    )["GENERATION_PROJECT"]
     n = len(projects)
     projects = projects[~projects["GENERATION_PROJECT"].isin(predetermined)]
     print(f"Removed {n - len(projects)} projects that were predetermined plants.")
     del predetermined
     # Determine the gen_techs where gen_is_variable is always True and gen_is_baseload is always False.
     # Grouping and summing works since summing Falses gives 0 but summing Trues gives >0.
     projects["gen_is_not_variable"] = ~projects["gen_is_variable"]
-    grouped_projects = projects.groupby("gen_tech", as_index=False)[["gen_is_not_variable", "gen_is_baseload"]].sum()
+    grouped_projects = projects.groupby("gen_tech", as_index=False)[
+        ["gen_is_not_variable", "gen_is_baseload"]
+    ].sum()
     grouped_projects = grouped_projects[
-        (grouped_projects["gen_is_not_variable"] == 0) & (grouped_projects["gen_is_baseload"] == 0)]
+        (grouped_projects["gen_is_not_variable"] == 0)
+        & (grouped_projects["gen_is_baseload"] == 0)
+    ]
     gen_tech = grouped_projects["gen_tech"]
     del grouped_projects
     print(f"Aggregating for gen_tech: {gen_tech.values}")
@@ -207,86 +222,113 @@ def create_capacity_factors():
     print(f"Removed {n - len(projects)} projects that aren't of allowed gen_tech.")
 
     # Calculate the gen_forced_outage_rate and verify it is identical for all the projects within the same group
-    outage_rates = projects.groupby("gen_tech", as_index=False)["gen_forced_outage_rate"]
+    outage_rates = projects.groupby("gen_tech", as_index=False)[
+        "gen_forced_outage_rate"
+    ]
     if (outage_rates.nunique()["gen_forced_outage_rate"] - 1).sum() != 0:
-        outage_rates = outage_rates.nunique().set_index("gen_tech")["gen_forced_outage_rate"] - 1
+        outage_rates = (
+            outage_rates.nunique().set_index("gen_tech")["gen_forced_outage_rate"] - 1
+        )
         outage_rates = outage_rates[outage_rates != 0]
         raise Exception(
-            f"These generation technologies have different forced outage rates: {outage_rates.index.values}")
-    outage_rates = outage_rates.mean()  # They're all the same so mean returns the proper value
+            f"These generation technologies have different forced outage rates: {outage_rates.index.values}"
+        )
+    outage_rates = (
+        outage_rates.mean()
+    )  # They're all the same so mean returns the proper value
     del projects
     print("Check passed: gen_forced_outage_rate is identical.")
 
     # Read the dispatch instructions
-    dispatch = pd.read_csv("outputs/dispatch.csv",
-                           usecols=["generation_project", "timestamp", "gen_tech", "gen_load_zone", "DispatchGen_MW",
-                                    "Curtailment_MW"],
-                           index_col=False,
-                           dtype={"generation_project": str})
+    dispatch = pd.read_csv(
+        "outputs/dispatch.csv",
+        usecols=[
+            "generation_project",
+            "timestamp",
+            "gen_tech",
+            "gen_load_zone",
+            "DispatchGen_MW",
+            "Curtailment_MW",
+        ],
+        index_col=False,
+        dtype={"generation_project": str},
+    )
     # Keep only valid projects
     dispatch = dispatch[dispatch["generation_project"].isin(valid_gens)]
     # Group by timestamp, gen_tech and load_zone
-    dispatch = dispatch.groupby(["timestamp", "gen_tech", "gen_load_zone"], as_index=False).sum()
+    dispatch = dispatch.groupby(
+        ["timestamp", "gen_tech", "gen_load_zone"], as_index=False
+    ).sum()
     # Get the DispatchUpperLimit from DispatchGen + Curtailment
-    dispatch["DispatchUpperLimit"] = dispatch["DispatchGen_MW"] + dispatch["Curtailment_MW"]
+    dispatch["DispatchUpperLimit"] = (
+        dispatch["DispatchGen_MW"] + dispatch["Curtailment_MW"]
+    )
     dispatch = dispatch.drop(["DispatchGen_MW", "Curtailment_MW"], axis=1)
 
     # Add the period to each row by merging with outputs/timestamp.csv
-    timestamps = pd.read_csv("outputs/timestamps.csv",
-                             usecols=["timestamp", "period"],
-                             index_col=False)
+    timestamps = pd.read_csv(
+        "outputs/timestamps.csv", usecols=["timestamp", "period"], index_col=False
+    )
     dispatch = dispatch.merge(
-        timestamps,
-        on="timestamp",
-        how='left',
-        validate="many_to_one"
+        timestamps, on="timestamp", how="left", validate="many_to_one"
     )
     del timestamps
 
     # Read the gen_cap.csv
-    cap = pd.read_csv("outputs/gen_cap.csv",
-                      usecols=["GENERATION_PROJECT", "PERIOD", "gen_tech", "gen_load_zone", "GenCapacity"],
-                      index_col=False,
-                      dtype={"GENERATION_PROJECT": str}).rename({"PERIOD": "period"}, axis=1)
+    cap = pd.read_csv(
+        "outputs/gen_cap.csv",
+        usecols=[
+            "GENERATION_PROJECT",
+            "PERIOD",
+            "gen_tech",
+            "gen_load_zone",
+            "GenCapacity",
+        ],
+        index_col=False,
+        dtype={"GENERATION_PROJECT": str},
+    ).rename({"PERIOD": "period"}, axis=1)
     # Keep only valid projects
-    cap = cap[cap["GENERATION_PROJECT"].isin(valid_gens)].drop("GENERATION_PROJECT", axis=1)
+    cap = cap[cap["GENERATION_PROJECT"].isin(valid_gens)].drop(
+        "GENERATION_PROJECT", axis=1
+    )
     # Sum for the tech, period and load zone
     cap = cap.groupby(["period", "gen_tech", "gen_load_zone"], as_index=False).sum()
     # Merge onto dispatch
     dispatch = dispatch.merge(
         cap,
         on=["period", "gen_tech", "gen_load_zone"],
         how="left",
-        validate="many_to_one"
+        validate="many_to_one",
     )
     del cap
 
     # Filter out zones with no buildout
     is_no_buildout = dispatch["GenCapacity"] == 0
-    missing_data = dispatch\
-        [is_no_buildout]\
-        [["period", "gen_tech", "gen_load_zone"]]\
-        .drop_duplicates()\
-        .groupby(["period", "gen_tech"], as_index=False)["gen_load_zone"]\
-        .nunique()\
+    missing_data = (
+        dispatch[is_no_buildout][["period", "gen_tech", "gen_load_zone"]]
+        .drop_duplicates()
+        .groupby(["period", "gen_tech"], as_index=False)["gen_load_zone"]
+        .nunique()
         .rename({"gen_load_zone": "Number of Load Zones"}, axis=1)
+    )
     if missing_data["Number of Load Zones"].sum() > 0:
         warnings.warn(
-            f"Unable to make capacity factors for the following categories since total capacity in those zones is 0.\n{missing_data}")
+            f"Unable to make capacity factors for the following categories since total capacity in those zones is 0.\n{missing_data}"
+        )
     dispatch = dispatch[~is_no_buildout]
 
     # Merge outage rates onto dispatch
-    dispatch = dispatch.merge(
-        outage_rates,
-        on="gen_tech"
-    )
+    dispatch = dispatch.merge(outage_rates, on="gen_tech")
     del outage_rates
 
     dispatch["gen_max_capacity_factor"] = dispatch["DispatchUpperLimit"] / (
-                dispatch["GenCapacity"] * (1 - dispatch["gen_forced_outage_rate"]))
-    dispatch = dispatch[["gen_tech", "gen_load_zone", "timestamp", "gen_max_capacity_factor"]]
+        dispatch["GenCapacity"] * (1 - dispatch["gen_forced_outage_rate"])
+    )
+    dispatch = dispatch[
+        ["gen_tech", "gen_load_zone", "timestamp", "gen_max_capacity_factor"]
+    ]
     dispatch.to_csv("zonal_capacity_factors.csv", index=False)
 
 
-if __name__=="__main__":
-    create_capacity_factors()
+if __name__ == "__main__":
+    create_capacity_factors()
