From 851098ead7c03d0184f1ea463ecacf4d60b3e887 Mon Sep 17 00:00:00 2001
From: Daniel Olsen <daniel@breakthroughenergy.org>
Date: Mon, 7 Mar 2022 09:33:43 -0800
Subject: [PATCH] refactor: remove query calls which are incompatible with
 numexpr

---
 .../gather/griddata/hifld/data_process/generators.py | 12 +++++++-----
 .../griddata/hifld/data_process/transmission.py      |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/prereise/gather/griddata/hifld/data_process/generators.py b/prereise/gather/griddata/hifld/data_process/generators.py
index ff9818c2e..0ac72514a 100644
--- a/prereise/gather/griddata/hifld/data_process/generators.py
+++ b/prereise/gather/griddata/hifld/data_process/generators.py
@@ -115,10 +115,12 @@ def quadratic(x, a, b, c):
     try:
         epa_key = crosswalk_translation.loc[eia_key].values[0]
         samples = epa_ampd_groupby.get_group(epa_key)
-        filtered_samples = samples.query(
-            f"(not `{x_col}`.isnull()) and `{x_col}` != 0 "
-            f"and (not `{y_col}`.isnull()) and `{y_col}` != 0"
-        )
+        filtered_samples = samples.loc[
+            (samples[x_col] != 0)
+            & (samples[y_col] != 0)
+            & ~samples[x_col].isnull()
+            & ~samples[y_col].isnull()
+        ]
         if len(filtered_samples[x_col].unique()) < min_unique_x:
             return default_return
         if len(filtered_samples) < min_points:
@@ -321,7 +323,7 @@ def build_plant(bus, substations, kwargs={}):
     filter_suspicious_heat_rates(generators)
     augment_missing_heat_rates(generators)
     # Drop generators whose heat rates can't be estimated
-    to_be_dropped = generators.query("h1.isnull()")
+    to_be_dropped = generators.loc[generators["h1"].isnull()]
     print(
         f"Dropping generators whose heat rates can't be estimated: {len(to_be_dropped)}"
         f" out of {len(generators)}, {round(to_be_dropped['Pmax'].sum(), 0)} MW out of "
diff --git a/prereise/gather/griddata/hifld/data_process/transmission.py b/prereise/gather/griddata/hifld/data_process/transmission.py
index 8a7da6ed3..142e6908f 100644
--- a/prereise/gather/griddata/hifld/data_process/transmission.py
+++ b/prereise/gather/griddata/hifld/data_process/transmission.py
@@ -408,7 +408,7 @@ def map_via_neighbor_voltages(lines, neighbors, func, method_name):
         :param str method_name: method name to print once no more updates can be found.
         """
         while True:
-            missing = lines.query("VOLTAGE.isnull()")
+            missing = lines.loc[lines["VOLTAGE"].isnull()]
             if len(missing) == 0:
                 print(f"No more missing voltages remain after neighbor {method_name}")
                 break