feat: add logic to map_generator_to_sub_by_location

prereise/gather/griddata/hifld/data_process/generators.py

@@ -1,18 +1,53 @@
 import pandas as pd
+from powersimdata.utility.distance import haversine
 
 from prereise.gather.griddata.hifld.data_access import load
 
 
-def map_generator_to_bus_by_location(generator, substation_groupby):
+def map_generator_to_sub_by_location(generator, substation_groupby):
+    """Determine a likely substation for a generator to be connected to. Priority order
+    of mapping is: 1) if location is available and one or more substations exist in that
+    ZIP code, map by location to closest substation within that ZIP code, 2) if location
+    is available but no substations exist in that ZIP code, map to the closest
+    substation within neighboring ZIP codes, 3) if only ZIP code is available
+    (no location), and one or more substations exist, map to an arbitrarily chosen
+    substation within that ZIP code, 4) if if only ZIP code is available (no location)
+    but no substations exist in that ZIP code, return NA.
+
+    :param pandas.Series generator: one generating unit from data frame.
+    :param pandas.GroupBy substation_groupby: data frame of substations, grouped by
+        (interconnect, ZIP).
+    """
+    lookup_params = tuple(generator.loc[["interconnect", "zip"]])
     if pd.isna(generator["lat"]) or pd.isna(generator["lon"]):
-        return pd.NA
-    if generator.loc["zip"] in substation_groupby:
-        # select the closest substation within that zip code
-        pass
-    else:
-        # select the closest substation
-        # (maybe filter based on a set of 'nearby' substations, to minimize computation)
-        pass
+        # No location available
+        try:
+            matching_subs = substation_groupby.get_group(lookup_params)
+            return matching_subs.index[0]
+        except KeyError:
+            return pd.NA
+    try:
+        # This ZIP code contains substations, this block will execute successfully
+        matching_subs = substation_groupby.get_group(lookup_params)
+    except KeyError:
+        # If this ZIP code does not contain substations, this block will execute, and
+        # we select a set of 'nearby' substations
+        zip_range = [generator.loc["zip"] + offset for offset in range(-100, 101)]
+        matching_subs = pd.concat(
+            [
+                substation_groupby.get_group((generator.loc["interconnect"], z))
+                for z in zip_range
+            ]
+        )
+    distance_to_subs = matching_subs.apply(
+        lambda x: haversine((x.lat, x.lon), (generator.lat, generator.lon)),
+        axis=1,
+    )
+    return distance_to_subs.idxmin()
+
+
+def map_generator_to_bus_by_sub(generator, bus, bus2sub):
+    pass
 
 
 def add_locations_to_generators(generators, power_plants):
@@ -94,8 +129,11 @@ def build_plant():
     # Add more information
     add_locations_to_generators(generators, power_plants)
     add_interconnects_to_generators(generators)
+    generators["sub_id"] = generators.apply(
+        lambda x: map_generator_to_sub_by_location(x, substation_groupby), axis=1
+    )
     generators["bus_id"] = generators.apply(
-        lambda x: map_generator_to_bus_by_location(x, substation_groupby), axis=1
+        lambda x: map_generator_to_bus_by_sub(x, bus, bus2sub), axis=1  # noqa: F821
     )
     generators["Pmax"] = generator_groupby["NAMPLT_CAP"].sum()
     add_pmin_to_generators(generators, eia_form_860)