refactor: simplify looping using pandas

prereise/gather/winddata/hrrr/calculations.py

@@ -81,27 +81,21 @@ def calculate_pout(wind_farms, start_dt, end_dt, directory):
         dt2    POUT        POUT
     """
 
-    wind_farm_ids = wind_farms.index
-    turbine_types = [
-        "Offshore"
-        if wind_farms.loc[i].type == "wind_offshore"
-        else id2abv[wind_farms.loc[i].zone_id]
-        for i in wind_farm_ids
-    ]
-    wind_farm_ct = len(wind_farms)
+    turbine_types = wind_farms.apply(
+        lambda x: "Offshore" if x["type"] == "wind_offshore" else id2abv[x["zone_id"]],
+        axis=1,
+    )
 
     turbine_power_curves = get_turbine_power_curves()
     state_power_curves = get_state_power_curves()
     wind_data_lat_long = get_wind_data_lat_long(start_dt, directory)
     wind_farm_to_closest_wind_grid_indices = find_closest_wind_grids(
         wind_farms, wind_data_lat_long
     )
-    dts = [
-        dt for dt in pd.date_range(start=start_dt, end=end_dt, freq="H").to_pydatetime()
-    ]
+    dts = pd.date_range(start=start_dt, end=end_dt, freq="H").to_pydatetime()
     # Fetch wind speed data for each wind farm (or store NaN as applicable)
-    wind_speed_data = np.empty((len(dts), len(wind_farm_to_closest_wind_grid_indices)))
-    for i, dt in tqdm(enumerate(dts)):
+    wind_speed_data = pd.DataFrame(index=dts, columns=wind_farms.index, dtype=float)
+    for dt in tqdm(dts):
         gribs = pygrib.open(formatted_filename(dt))
         try:
             u_component = gribs.select(name=U_COMPONENT_SELECTOR)[0].values.flatten()
@@ -112,13 +106,13 @@ def calculate_pout(wind_farms, start_dt, end_dt, directory):
             wind_farm_specific_v_component = v_component[
                 wind_farm_to_closest_wind_grid_indices
             ]
-            wind_speed_data[i] = np.sqrt(
+            wind_speed_data.loc[dt] = np.sqrt(
                 pow(wind_farm_specific_u_component, 2)
                 + pow(wind_farm_specific_v_component, 2)
             )
         except ValueError:
             # If the GRIB file is empty, no wind speed values can be selected
-            wind_speed_data[i] = np.nan
+            wind_speed_data.loc[dt] = np.nan
 
     # For each column, linearly interpolate any NaN values
     linear(wind_speed_data)
@@ -128,12 +122,12 @@ def calculate_pout(wind_farms, start_dt, end_dt, directory):
             get_power(
                 turbine_power_curves,
                 state_power_curves,
-                wind_speed_data[i, j],
-                turbine_types[j],
+                wind_speed_data.loc[dt, w],
+                turbine_types.loc[w],
             )
-            for j in range(wind_farm_ct)
+            for w in wind_farms.index
         ]
-        for i, _ in tqdm(enumerate(dts))
+        for dt in tqdm(dts)
     ]
     df = pd.DataFrame(data=wind_power_data, index=dts, columns=wind_farms.index)