fix: add sub and bus2sub entries, as needed, for new buses

powersimdata/input/transform_grid.py

@@ -231,6 +231,7 @@ def _add_bus(self):
         }
         new_bus = {c: 0 for c in bus.columns}
         for entry in self.ct["new_bus"]:
+            # Add to the bus dataframe
             new_bus["type"] = 1
             new_bus["Pd"] = entry["Pd"]
             new_bus["zone_id"] = entry["zone_id"]
@@ -239,13 +240,46 @@ def _add_bus(self):
             new_bus["loss_zone"] = 1
             new_bus["Vmax"] = 1.1
             new_bus["Vmin"] = 0.9
-            new_bus["interconnect"] = zone2interconnect[entry["zone_id"]]
-            new_bus["lat"] = entry["lat"]
-            new_bus["lon"] = entry["lon"]
-            new_index = [self.grid.bus.index.max() + 1]
+            interconnect = zone2interconnect[entry["zone_id"]]
+            new_bus["interconnect"] = interconnect
+            lat, lon = entry["lat"], entry["lon"]
+            new_bus["lat"] = lat
+            new_bus["lon"] = lon
+            new_bus_id = [self.grid.bus.index.max() + 1]
             self.grid.bus = self.grid.bus.append(
-                pd.DataFrame(new_bus, index=new_index), sort=False
+                pd.DataFrame(new_bus, index=new_bus_id), sort=False
             )
+            # Add to substation & bus2sub mapping dataframes
+            sub = self.grid.sub
+            if (lat, lon) in sub.groupby(["lat", "lon"]).groups.keys():
+                # If there are multiple matching substations, arbitrarily grab the first
+                sub_id = sub[(sub.lat == lat) & (sub.lon == lon)].index[0]
+                new_row = pd.DataFrame(
+                    {"sub_id": sub_id, "interconnect": interconnect},
+                    index=[new_bus_id],
+                )
+                self.grid.bus2sub = self.grid.bus2sub.append(new_row, sort=False)
+            else:
+                # Create a new substation
+                new_sub_id = sub.index.max() + 1
+                interconnect_sub = sub[sub.interconnect == interconnect]
+                new_interconnect_sub_id = interconnect_sub.interconnect_sub_id.max() + 1
+                new_row = pd.DataFrame(
+                    {"sub_id": new_sub_id, "interconnect": interconnect},
+                    index=[new_bus_id],
+                )
+                self.grid.bus2sub = self.grid.bus2sub.append(new_row, sort=False)
+                new_row = pd.DataFrame(
+                    {
+                        "name": f"NEW {new_sub_id}",
+                        "interconnect_sub_id": new_interconnect_sub_id,
+                        "lat": lat,
+                        "lon": lon,
+                        "interconnect": interconnect,
+                    },
+                    index=[new_sub_id],
+                )
+                self.grid.sub = sub.append(new_row, sort=False)
 
     def _add_dcline(self):
         """Adds HVDC line(s) to the grid"""