Merge pull request #94 from Breakthrough-Energy/bainan/recover_indices

feat: recover plant and branch indices

switchwrapper/grid_to_switch.py

@@ -4,7 +4,7 @@
 from haversine import haversine
 
 from switchwrapper import const
-from switchwrapper.helpers import make_indices
+from switchwrapper.helpers import make_branch_indices, make_plant_indices
 
 
 def grid_to_switch(grid, output_folder):
@@ -261,7 +261,7 @@ def build_generation_projects_info(plant, single_segment_slope, average_fuel_cos
     :return: (*pandas.DataFrame*) -- data frame of generation project info.
     """
     # Extract information from inputs
-    original_plant_indices, hypothetical_plant_indices = make_indices(plant.index)
+    original_plant_indices, hypothetical_plant_indices = make_plant_indices(plant.index)
     all_plant_indices = original_plant_indices + hypothetical_plant_indices
 
     # Use inputs for intermediate calculations
@@ -318,7 +318,7 @@ def build_gen_build_costs(plant, cost_at_min_power, inv_period):
     :return: (*pandas.DataFrame*) -- data frame of existing and hypothetical generators.
     """
     # Build lists for each columns, which apply to one year
-    original_plant_indices, hypothetical_plant_indices = make_indices(plant.index)
+    original_plant_indices, hypothetical_plant_indices = make_plant_indices(plant.index)
     overnight_costs = plant["type"].map(const.investment_costs_by_type).tolist()
     gen_fixed_om = (cost_at_min_power / plant.Pmax).fillna(0.0).tolist()
 
@@ -359,7 +359,7 @@ def build_gen_build_predetermined(plant):
         },
         inplace=True,
     )
-    original_plant_indices, _ = make_indices(plant.index)
+    original_plant_indices, _ = make_plant_indices(plant.index)
     gen_build_predetermined["GENERATION_PROJECT"] = original_plant_indices
     gen_build_predetermined = gen_build_predetermined[
         ["GENERATION_PROJECT", "build_year", "gen_predetermined_cap"]
@@ -444,7 +444,7 @@ def build_aclines(grid):
             grid.bus.loc[acline["to_bus_id"], "baseKV"],
         )
     )
-    acline["branch_id"] = acline["branch_id"].apply(lambda x: str(x) + "ac")
+    acline["branch_id"] = make_branch_indices(acline["branch_id"])
     return acline.round(2)
 
 
@@ -464,7 +464,7 @@ def build_dclines(grid):
         )
     )
     dcline["trans_efficiency"] = 0.99
-    dcline["dcline_id"] = dcline["dcline_id"].apply(lambda x: str(x) + "dc")
+    dcline["dcline_id"] = make_branch_indices(dcline["dcline_id"], dc=True)
     dcline.rename(columns={"dcline_id": "branch_id", "Pmax": "rateA"}, inplace=True)
     return dcline.round(2)
 

switchwrapper/helpers.py

@@ -3,18 +3,6 @@
 import pandas as pd
 
 
-def make_indices(plant_ids):
-    """Make the indices for existing and hypothetical generators for input to Switch.
-
-    :param iterable plant_ids: plant IDs.
-    :return: (*tuple*) -- The first element is a list of indices for existing generators
-        and the second element is a list of indices for hypothetical generators.
-    """
-    original_plant_indices = [f"g{p}" for p in plant_ids]
-    hypothetical_plant_indices = [f"{o}i" for o in original_plant_indices]
-    return original_plant_indices, hypothetical_plant_indices
-
-
 def match_variables(variables, pattern, columns):
     """Search through dictionary of variables, extracting data frame of values.
 
@@ -37,3 +25,84 @@ def match_variables(variables, pattern, columns):
         ]
     )
     return df
+
+
+def make_plant_indices(plant_ids):
+    """Make the indices for existing and hypothetical generators for input to Switch.
+
+    :param iterable plant_ids: plant IDs.
+    :return: (*tuple*) -- The first element is a list of indices for existing generators
+        and the second element is a list of indices for hypothetical generators.
+    """
+    original_plant_indices = [f"g{p}" for p in plant_ids]
+    hypothetical_plant_indices = [f"{o}i" for o in original_plant_indices]
+    return original_plant_indices, hypothetical_plant_indices
+
+
+def make_branch_indices(branch_ids, dc=False):
+    """Make the indices of existing branch for input to Switch.
+
+    :param iterable branch_ids: list of original branch ids.
+    :param bool dc: branch_ids are for dclines or not, defaults to False.
+    :return: (*list*) -- list of branch indices for input to Switch
+    """
+    return [f"{i}dc" if dc else f"{i}ac" for i in branch_ids]
+
+
+def recover_plant_indices(switch_plant_ids):
+    """Recover the plant indices from Switch outputs.
+
+    :param iterable switch_plant_ids: Switch plant indices.
+    :return: (*pandas.Series*) -- indices are original plant ids with new plants
+        added, values are Switch plant indices.
+    """
+    plant_ids = dict()
+    for ind in switch_plant_ids[::-1]:
+        if ind[-1] != "i":
+            last_original_plant_id = int(ind[1:])
+            break
+    cnt = 0
+    for ind in switch_plant_ids:
+        if ind[-1] != "i":
+            plant_ids[int(ind[1:])] = ind
+        else:
+            cnt += 1
+            plant_ids[last_original_plant_id + cnt] = ind
+    return pd.Series(plant_ids)
+
+
+def recover_branch_indices(switch_branch_ids):
+    """Recover the branch indices from Switch outputs.
+
+    :param iterable switch_branch_ids: Switch branch indices.
+    :return: (*tuple*) -- a pair of pandas.Series objects for acline and dcline
+        respectively, which are indexed by original branch ids and values are
+        corresponding Switch branch indices.
+    """
+    ac_branch_ids = dict()
+    dc_branch_ids = dict()
+    for ind in switch_branch_ids:
+        if ind[-2:] == "ac":
+            ac_branch_ids[int(ind[:-2])] = ind
+        else:
+            dc_branch_ids[int(ind[:-2])] = ind
+    return pd.Series(ac_branch_ids), pd.Series(dc_branch_ids, dtype=str)
+
+
+def branch_indices_to_bus_tuple(grid):
+    """Map the branch indices to from/to bus tuples based on a grid instance.
+
+    :param powersimdata.input.grid.Grid grid: grid instance.
+    :return: (*tuple*) -- a pair of pandas.Series objects for acline and dcline
+        respectively, which are indexed by original branch ids and values are
+        corresponding tuples (from_bus_id, to_bus_id).
+    """
+    acline = pd.Series(
+        list(zip(grid.branch["from_bus_id"], grid.branch["to_bus_id"])),
+        index=grid.branch.index,
+    )
+    dcline = pd.Series(
+        list(zip(grid.dcline["from_bus_id"], grid.dcline["to_bus_id"])),
+        index=grid.dcline.index,
+    )
+    return acline, dcline

switchwrapper/tests/test_helpers.py

@@ -0,0 +1,60 @@
+import pandas as pd
+from powersimdata.tests.mock_grid import MockGrid
+
+from switchwrapper.helpers import (
+    branch_indices_to_bus_tuple,
+    recover_branch_indices,
+    recover_plant_indices,
+)
+
+mock_branch = {
+    "branch_id": [101, 102, 103],
+    "from_bus_id": [1, 2, 3],
+    "to_bus_id": [2, 3, 1],
+}
+
+mock_dcline = {
+    "dcline_id": [0],
+    "from_bus_id": [1],
+    "to_bus_id": [2],
+}
+
+
+def test_recover_plant_indices():
+    args = [
+        ["g1", "g2", "g1i", "g2i"],
+        ["g1", "g2", "g1i"],
+        ["g1", "g2"],
+    ]
+    expected_return = [
+        pd.Series({1: "g1", 2: "g2", 3: "g1i", 4: "g2i"}),
+        pd.Series({1: "g1", 2: "g2", 3: "g1i"}),
+        pd.Series({1: "g1", 2: "g2"}),
+    ]
+    for a, e in zip(args, expected_return):
+        assert e.equals(recover_plant_indices(a))
+
+
+def test_recover_branch_indices():
+    args = [["1ac", "2ac", "0dc"], ["1ac", "2ac"]]
+    expected_return = [
+        (pd.Series({1: "1ac", 2: "2ac"}), pd.Series({0: "0dc"})),
+        (pd.Series({1: "1ac", 2: "2ac"}), pd.Series(dtype=str)),
+    ]
+    for a, e in zip(args, expected_return):
+        assert all([s.equals(e[i]) for i, s in enumerate(recover_branch_indices(a))])
+
+
+def test_branch_indices_to_bus_tuple():
+    grid = MockGrid(grid_attrs={"branch": mock_branch, "dcline": mock_dcline})
+    acline, dcline = branch_indices_to_bus_tuple(grid)
+    expected_acline = pd.Series(
+        {
+            101: (1, 2),
+            102: (2, 3),
+            103: (3, 1),
+        }
+    )
+    expected_dcline = pd.Series({0: (1, 2)})
+    assert expected_acline.equals(acline)
+    assert expected_dcline.equals(dcline)