Simplify adjust_no_electric_network by not creating new buses

etrago/cluster/electrical.py

@@ -137,16 +137,9 @@ def adjust_no_electric_network(etrago, busmap, cluster_met):
         "rural_heat": "rural_heat_pump",
     }
 
-    # no_elec_to_cluster maps the no electrical buses to the eHV/kmean bus
-    no_elec_to_cluster = pd.DataFrame(
-        columns=["cluster", "carrier", "new_bus"]
-    ).set_index("new_bus")
-
-    max_bus = max([int(item) for item in network.buses.index.to_list()])
-
     no_elec_conex = []
-    # busmap2 maps all the no electrical buses to the new buses based on the
-    # eHV network
+    # busmap2 defines how the no electrical buses directly connected to AC
+    # are going to be clustered
     busmap2 = {}
 
     # Map crossborder AC buses in case that they were not part of the k-mean clustering
@@ -161,52 +154,26 @@ def adjust_no_electric_network(etrago, busmap, cluster_met):
         for bus_out in ac_buses_out.index:
             busmap2[bus_out] = bus_out
 
+    busmap3 = pd.DataFrame(columns=["elec_bus", "carrier", "cluster"])
     for bus_ne in network2.buses.index:
-        bus_hv = -1
         carry = network2.buses.loc[bus_ne, "carrier"]
-
-        if (
-            len(
-                network2.links[
-                    (network2.links["bus1"] == bus_ne)
-                    & (network2.links["carrier"] == map_carrier[carry])
-                ]
-            )
-            > 0
-        ):
+        busmap3.at[bus_ne, "carrier"] = carry
+        try:
             df = network2.links[
                 (network2.links["bus1"] == bus_ne)
                 & (network2.links["carrier"] == map_carrier[carry])
             ].copy()
             df["elec"] = df["bus0"].isin(busmap.keys())
-            df = df[df["elec"] == True]
-            if len(df) > 0:
-                bus_hv = df["bus0"][0]
-
-        if bus_hv == -1:
-            busmap2[bus_ne] = str(bus_ne)
+            bus_hv = df[df["elec"] == True]["bus0"][0]
+            busmap3.at[bus_ne, "elec_bus"] = busmap[bus_hv]
+        except:
             no_elec_conex.append(bus_ne)
-            continue
-
-        if (
-            (no_elec_to_cluster.cluster == busmap[bus_hv])
-            & (no_elec_to_cluster.carrier == carry)
-        ).any():
-            bus_cluster = no_elec_to_cluster[
-                (no_elec_to_cluster.cluster == busmap[bus_hv])
-                & (no_elec_to_cluster.carrier == carry)
-            ].index[0]
-        else:
-            bus_cluster = str(max_bus + 1)
-            max_bus = max_bus + 1
-            new = pd.DataFrame(
-                {"cluster": busmap[bus_hv], "carrier": carry},
-                index=[bus_cluster],
-            )
+            busmap3.at[bus_ne, "elec_bus"] = bus_ne
 
-            no_elec_to_cluster = pd.concat([no_elec_to_cluster, new])
+    for a, df in busmap3.groupby(["elec_bus", "carrier"]):
+        busmap3.loc[df.index, "cluster"] = df.index[0]
 
-        busmap2[bus_ne] = bus_cluster
+    busmap3 = busmap3["cluster"].to_dict()
 
     if no_elec_conex:
         logger.info(
@@ -215,17 +182,17 @@ def adjust_no_electric_network(etrago, busmap, cluster_met):
         )
 
     # rural_heat_store buses are clustered based on the AC buses connected to
-    # their corresponding rural_heat buses
+    # their corresponding rural_heat buses. Results saved in busmap4
     links_rural_store = etrago.network.links[
         etrago.network.links.carrier == "rural_heat_store_charger"
     ].copy()
 
-    busmap3 = {}
-    links_rural_store["to_ac"] = links_rural_store["bus0"].map(busmap2)
+    busmap4 = {}
+    links_rural_store["to_ac"] = links_rural_store["bus0"].map(busmap3)
     for rural_heat_bus, df in links_rural_store.groupby("to_ac"):
         cluster_bus = df.bus1.iat[0]
         for rural_store_bus in df.bus1:
-            busmap3[rural_store_bus] = cluster_bus
+            busmap4[rural_store_bus] = cluster_bus
 
     # Add the gas buses to the busmap and map them to themself
     for gas_bus in network.buses[
@@ -236,28 +203,8 @@ def adjust_no_electric_network(etrago, busmap, cluster_met):
     ].index:
         busmap2[gas_bus] = gas_bus
 
-    busmap = {**busmap, **busmap2, **busmap3}
-
-    # The new buses based on the eHV network for not electrical buses are
-    # created
-    if cluster_met in ["kmeans", "kmedoids-dijkstra"]:
-        network.madd(
-            "Bus",
-            names=no_elec_to_cluster.index,
-            carrier=no_elec_to_cluster.carrier,
-        )
+    busmap = {**busmap, **busmap2, **busmap3, **busmap4}
 
-    else:
-        network.madd(
-            "Bus",
-            names=no_elec_to_cluster.index,
-            carrier=no_elec_to_cluster.carrier.values,
-            x=network.buses.loc[no_elec_to_cluster.cluster.values, "x"].values,
-            y=network.buses.loc[no_elec_to_cluster.cluster.values, "y"].values,
-            country=network.buses.loc[
-                no_elec_to_cluster.cluster.values, "country"
-            ].values,
-        )
     return network, busmap