refactor: improve efficiency of _add_bus

powersimdata/input/transform_grid.py

@@ -229,6 +229,7 @@ def _add_bus(self):
         zone2interconnect = {
             k: v[0] for k, v in bus.groupby("zone_id").interconnect.unique().items()
         }
+        latlon2sub = self.grid.sub.groupby(["lat", "lon"]).groups
         new_bus = {c: 0 for c in bus.columns}
         for entry in self.ct["new_bus"]:
             # Add to the bus dataframe
@@ -250,17 +251,17 @@ def _add_bus(self):
                 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 (lat, lon) in latlon2sub:
                 # If there are multiple matching substations, arbitrarily grab the first
-                sub_id = sub[(sub.lat == lat) & (sub.lon == lon)].index[0]
+                sub_id = latlon2sub[(lat, lon)][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
+                sub = self.grid.sub
                 new_sub_id = sub.index.max() + 1
                 interconnect_sub = sub[sub.interconnect == interconnect]
                 new_interconnect_sub_id = interconnect_sub.interconnect_sub_id.max() + 1
@@ -280,6 +281,7 @@ def _add_bus(self):
                     index=[new_sub_id],
                 )
                 self.grid.sub = sub.append(new_row, sort=False)
+                latlon2sub[(lat, lon)] = new_sub_id
 
     def _add_dcline(self):
         """Adds HVDC line(s) to the grid"""