Integration Test fixes

quartz_solar_forecast/eval/pv.py

@@ -1,93 +1,130 @@
+import glob
 import os
 
 import numpy as np
 import pandas as pd
-import xarray as xr
+import pyarrow as pa
+import pyarrow.dataset as ds
 from huggingface_hub import HfFileSystem
 
 fs = HfFileSystem()
 
 
-def get_pv_metadata(testset: pd.DataFrame):
-    # download from hugginface or load from cache
+def get_pv_metadata(testset: pd.DataFrame) -> pd.DataFrame:
+    """Merge metadata (lat/lon/capacity) with testset of pv_id + timestamp."""
     cache_dir = "data/pv"
     metadata_file = f"{cache_dir}/metadata.csv"
-    if not os.path.exists(metadata_file):
+
+    if not os.path.exists(metadata_file) or os.path.getsize(metadata_file) == 0:
         os.makedirs(cache_dir, exist_ok=True)
         fs.get("datasets/openclimatefix/uk_pv/metadata.csv", metadata_file)
 
-    # Load in the dataset
     metadata_df = pd.read_csv(metadata_file)
 
-    # join metadata with testset
+    # align schema
     metadata_df = metadata_df.rename(columns={"ss_id": "pv_id"})
+
     combined_data = testset.merge(metadata_df, on="pv_id", how="left")
 
-    # only keep the columns we need
+    # keep only useful columns
     combined_data = combined_data[
-        ["pv_id", "timestamp", "latitude_rounded", "longitude_rounded", "kwp"]
-    ]
-
-    # rename latitude_rounded to latitude and longitude_rounded to longitude
-    combined_data = combined_data.rename(
+        ["pv_id", "timestamp", "latitude_rounded", "longitude_rounded", "kWp"]
+    ].rename(
         columns={
             "latitude_rounded": "latitude",
             "longitude_rounded": "longitude",
-            "kwp": "capacity",
+            "kWp": "capacity",
         }
     )
 
-    # format datetime
     combined_data["timestamp"] = pd.to_datetime(combined_data["timestamp"])
-
     return combined_data
 
 
-def get_pv_truth(testset: pd.DataFrame):
-    print("Loading PV data")
+FOLDER_TO_TIME_RES = {
+    "5_minutely": "5min",
+    "30_minutely": "30min",
+}
+
+
+def get_pv_truth(
+    testset: pd.DataFrame, horizon_hours: int = 48, folder_name: str = "30_minutely"
+) -> pd.DataFrame:
+    """
+    Fetch PV generation truth values for given testset.
+    Optimized for performance using Arrow predicate filtering.
+    """
 
-    # download from hugginface or load from cache
     cache_dir = "data/pv"
-    metadata_file = f"{cache_dir}/pv.netcdf"
-    if not os.path.exists(metadata_file):
-        print("Loading from HF)")
+    parquet_dir = f"{cache_dir}/{folder_name}"
+
+    if not os.path.exists(parquet_dir):
+        print("Downloading PV parquet data from HuggingFace...")
         os.makedirs(cache_dir, exist_ok=True)
-        fs.get("datasets/openclimatefix/uk_pv/pv.netcdf", metadata_file)
-
-    # Load in the dataset
-    pv_ds = xr.open_dataset(metadata_file, engine="h5netcdf")
-
-    combined_data = []
-    for index, row in testset.iterrows():
-        print(f"Processing {index} of {len(testset)}")
-        pv_id = str(row["pv_id"])
-        base_datetime = pd.to_datetime(row["timestamp"])
-
-        # Calculate future timestamps up to the max horizon
-        for i in range(0, 49):  # 48 hours in steps of 1 hour
-            future_datetime = base_datetime + pd.DateOffset(hours=i)
-            horizon = i  # horizon in hours
-
-            try:
-                # Attempt to select data for the future datetime
-                selected_data = pv_ds[pv_id].sel(datetime=future_datetime)
-                value = selected_data.values.item()
-                value = value / 1000  # to convert from w to kw
-            except KeyError:
-                # If data is not found for the future datetime, set value as NaN
-                value = np.nan
-
-            # Add the data to the DataFrame
-            combined_data.append(
-                pd.DataFrame(
-                    {
-                        "pv_id": pv_id,
-                        "timestamp": future_datetime,
-                        "value": value,
-                        "horizon_hour": horizon,
-                    },
-                    index=[i],
-                )
-            )
-    combined_data = pd.concat(combined_data)
-    return combined_data
+        fs.get(f"datasets/openclimatefix/uk_pv/{folder_name}", cache_dir, recursive=True)
+
+    # Find all non-empty parquet files
+    files = glob.glob(f"{parquet_dir}/**/*.parquet", recursive=True)
+    non_empty_files = [f for f in files if os.path.getsize(f) > 0]
+
+    if not non_empty_files:
+        raise FileNotFoundError("No valid parquet files found (all are empty).")
+
+    # Prepare filtering parameters
+    unique_pv_ids = testset["pv_id"].unique().tolist()
+    min_time = pd.to_datetime(testset["timestamp"]).min()
+    max_time = min_time + pd.Timedelta(hours=horizon_hours)
+
+    # Ensure timestamps are timezone-aware
+    testset["timestamp"] = pd.to_datetime(testset["timestamp"], utc=True)
+
+    # Define dataset with filtering
+    dataset = ds.dataset(non_empty_files, format="parquet")
+
+    arrow_min_time = pa.scalar(min_time, type=pa.timestamp("ns", tz="UTC"))
+    arrow_max_time = pa.scalar(max_time, type=pa.timestamp("ns", tz="UTC"))
+
+    filter_expr = (
+        (ds.field("ss_id").isin(unique_pv_ids))
+        & (ds.field("datetime_GMT") >= arrow_min_time)
+        & (ds.field("datetime_GMT") <= arrow_max_time)
+    )
+
+    # Load filtered data only
+    table = dataset.to_table(filter=filter_expr)
+    pv_data = table.to_pandas()
+
+    # Ensure datetime column is parsed and aligned
+    pv_data["datetime_GMT"] = pd.to_datetime(pv_data["datetime_GMT"], utc=True)
+    time_resolution = FOLDER_TO_TIME_RES.get(folder_name)
+    if time_resolution is None:
+        raise ValueError(
+            f"Unknown folder_name '{folder_name}'. Please add it to FOLDER_TO_TIME_RES mapping."
+        )
+
+    pv_data["datetime_GMT"] = pv_data["datetime_GMT"].dt.floor(time_resolution)
+
+    # Expand testset for all horizons
+    horizons = np.arange(horizon_hours + 1)
+    expanded = testset.loc[testset.index.repeat(len(horizons))].copy()
+    expanded["horizon_hour"] = np.tile(horizons, len(testset))
+
+    # Calculate actual timestamp for each horizon
+    expanded["timestamp"] = expanded["timestamp"] + pd.to_timedelta(
+        expanded["horizon_hour"], unit="h"
+    )
+    expanded["timestamp"] = expanded["timestamp"].dt.floor(time_resolution)
+    # Merge
+    merged = expanded.merge(
+        pv_data,
+        left_on=["pv_id", "timestamp"],
+        right_on=["ss_id", "datetime_GMT"],
+        how="left",
+    )
+
+    # Convert to kWh
+    merged["value"] = merged["generation_Wh"] / 1000.0
+
+    result = merged[["pv_id", "timestamp", "value", "horizon_hour"]].copy()
+
+    return result

tests/integration/eval/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+
+@pytest.mark.parametrize("folder_name", ["30_minutely", "5_minutely"])
+def test_folder_processing(folder_name):
+    """
+    Run tests for both folder types using parameterization.
+    """
+    print(f"Running tests for folder: {folder_name}")
+    assert folder_name in ["30_minutely", "5_minutely"]

tests/integration/eval/test_pv.py

@@ -2,8 +2,8 @@
 import pandas as pd
 import pytest
 
-@pytest.mark.skip(reason="HF files have been changes"
-" - https://github.com/openclimatefix/open-source-quartz-solar-forecast/issues/292")
+# @pytest.mark.skip(reason="HF files have been changes"
+# " - https://github.com/openclimatefix/open-source-quartz-solar-forecast/issues/292")
 @pytest.mark.integration
 def test_get_pv_metadata():
     test_set_df = pd.DataFrame(
@@ -19,10 +19,11 @@ def test_get_pv_metadata():
     assert "latitude" in metadata_df.columns
 
 
-@pytest.mark.skip(reason="HF files have been changes"
-" - https://github.com/openclimatefix/open-source-quartz-solar-forecast/issues/292")
+# @pytest.mark.skip(reason="HF files have been changes"
+# " - https://github.com/openclimatefix/open-source-quartz-solar-forecast/issues/292")
+@pytest.mark.parametrize("folder_name", ["30_minutely", "5_minutely"])
 @pytest.mark.integration
-def test_get_pv():
+def test_get_pv(folder_name):
     # make test dataset file
     test_set_df = pd.DataFrame(
         [
@@ -34,4 +35,4 @@ def test_get_pv():
     )
 
     # Collect NWP data from Hugging Face, ICON. (Peter)
-    _ = get_pv_truth(test_set_df)
+    _ = get_pv_truth(test_set_df,folder_name=folder_name)