fix test_eval_forecast

openclimatefix · Dec 12, 2023 · 674e69f · 674e69f
1 parent d9f1adb
commit 674e69f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/quartz_solar_forecast/eval/forecast.py b/quartz_solar_forecast/eval/forecast.py
@@ -55,9 +55,10 @@ def run_forecast(pv_df: pd.DataFrame, nwp_df: pd.DataFrame, nwp_source="ICON") -
         # format
         times = nwp_site_df["time"]
         step = times - ts
-        nwp_site_df = nwp_site_df.drop(
-            columns=["timestamp", "latitude", "longitude", "pv_id"]
-        )
+        for c in ["timestamp", "latitude", "longitude", "pv_id"]:
+            if c in nwp_site_df.columns:
+                nwp_site_df = nwp_site_df.drop(columns=c)
+
         nwp_site_df.set_index("time", inplace=True, drop=True)
 
         if isinstance(ts, str):

diff --git a/tests/eval/test_eval_forecast.py b/tests/eval/test_eval_forecast.py
@@ -18,7 +18,7 @@ def test_run_forecast():
 
     nwp_df = pd.DataFrame(
         columns=[
-            "id",
+            "pv_id",
             "timestamp",
             "t",
             "prate",