""" A script to run backtest for PVNet for specific sites Use: - This script uses hydra to construct the config, just like in `run.py`. So you need to make sure that the data config is set up appropriate for the model being run in this script - The PVNet model checkpoint; the time range over which to make predictions are made; the site ids to produce forecasts for and the output directory where the results near the top of the script as hard coded user variables. These should be changed. ``` python scripts/backtest_sites.py ``` """ try: import torch.multiprocessing as mp mp.set_start_method("spawn", force=True) mp.set_sharing_strategy("file_system") except RuntimeError: pass import json import logging import os import sys import hydra import numpy as np import pandas as pd import torch import xarray as xr from huggingface_hub import hf_hub_download from huggingface_hub.constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME from ocf_data_sampler.sample.base import batch_to_tensor, copy_batch_to_device from ocf_datapipes.batch import ( BatchKey, NumpyBatch, stack_np_examples_into_batch, ) from ocf_datapipes.config.load import load_yaml_configuration from ocf_datapipes.load.pv.pv import OpenPVFromNetCDFIterDataPipe from ocf_datapipes.training.common import create_t0_and_loc_datapipes from ocf_datapipes.training.pvnet_site import ( DictDatasetIterDataPipe, _get_datapipes_dict, construct_sliced_data_pipeline, split_dataset_dict_dp, ) from ocf_datapipes.utils.consts import ELEVATION_MEAN, ELEVATION_STD from omegaconf import DictConfig from torch.utils.data import DataLoader, IterDataPipe, functional_datapipe from torch.utils.data.datapipes.iter import IterableWrapper from tqdm import tqdm from pvnet.load_model import get_model_from_checkpoints from pvnet.utils import SiteLocationLookup # ------------------------------------------------------------------ # USER CONFIGURED VARIABLES TO RUN THE SCRIPT # Directory path to save results output_dir = "PLACEHOLDER" # Local directory to load the PVNet checkpoint from. By default this should pull the best performing # checkpoint on the val set model_chckpoint_dir = "PLACEHOLDER" hf_revision = None hf_token = None hf_model_id = None # Forecasts will be made for all available init times between these start_datetime = "2022-05-08 00:00" end_datetime = "2022-05-08 00:30" # ------------------------------------------------------------------ # SET UP LOGGING logger = logging.getLogger(__name__) logging.basicConfig(stream=sys.stdout, level=logging.INFO) # ------------------------------------------------------------------ # DERIVED VARIABLES # This will run on GPU if it exists device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ------------------------------------------------------------------ # GLOBAL VARIABLES # The frequency of the pv site data FREQ_MINS = 30 # When sun as elevation below this, the forecast is set to zero MIN_DAY_ELEVATION = 0 # Add all pv site ids here that you wish to produce forecasts for ALL_SITE_IDS = [] # Need to be in ascending order ALL_SITE_IDS.sort() # ------------------------------------------------------------------ # FUNCTIONS @functional_datapipe("pad_forward_pv") class PadForwardPVIterDataPipe(IterDataPipe): """ Pads forecast pv. Sun position is calculated based off of pv time index and for t0's close to end of pv data can have wrong shape as pv starts to run out of data to slice for the forecast part. """ def __init__( self, pv_dp: IterDataPipe, forecast_duration: np.timedelta64, history_duration: np.timedelta64, time_resolution_minutes: np.timedelta64, ): """Init""" super().__init__() self.pv_dp = pv_dp self.forecast_duration = forecast_duration self.history_duration = history_duration self.time_resolution_minutes = time_resolution_minutes self.min_seq_length = history_duration // time_resolution_minutes def __iter__(self): """Iter""" for xr_data in self.pv_dp: t_end = ( xr_data.time_utc.data[0] + self.history_duration + self.forecast_duration + self.time_resolution_minutes ) time_idx = np.arange(xr_data.time_utc.data[0], t_end, self.time_resolution_minutes) if len(xr_data.time_utc.data) < self.min_seq_length: raise ValueError("Not enough PV data to predict") yield xr_data.reindex(time_utc=time_idx, fill_value=-1) def load_model_from_hf(model_id: str, revision: str, token: str): """ Loads model from HuggingFace """ model_file = hf_hub_download( repo_id=model_id, filename=PYTORCH_WEIGHTS_NAME, revision=revision, token=token, ) # load config file config_file = hf_hub_download( repo_id=model_id, filename=CONFIG_NAME, revision=revision, token=token, ) with open(config_file, "r", encoding="utf-8") as f: config = json.load(f) model = hydra.utils.instantiate(config) state_dict = torch.load(model_file, map_location=torch.device("cuda")) model.load_state_dict(state_dict) # type: ignore model.eval() # type: ignore return model def preds_to_dataarray(preds, model, valid_times, site_ids): """Put numpy array of predictions into a dataarray""" if model.use_quantile_regression: output_labels = [f"forecast_mw_plevel_{int(q*100):02}" for q in model.output_quantiles] output_labels[output_labels.index("forecast_mw_plevel_50")] = "forecast_mw" else: output_labels = ["forecast_mw"] preds = preds[..., np.newaxis] da = xr.DataArray( data=preds, dims=["pv_system_id", "target_datetime_utc", "output_label"], coords=dict( pv_system_id=site_ids, target_datetime_utc=valid_times, output_label=output_labels, ), ) return da # TODO change this to load the PV sites data (metadata?) def get_sites_ds(config_path: str) -> xr.Dataset: """Load site data from the path in the data config. Args: config_path: Path to the data configuration file Returns: xarray.Dataset of PVLive truths and capacities """ config = load_yaml_configuration(config_path) site_datapipe = OpenPVFromNetCDFIterDataPipe(pv=config.input_data.pv) ds_sites = next(iter(site_datapipe)) return ds_sites def get_available_t0_times(start_datetime, end_datetime, config_path): """Filter a list of t0 init-times to those for which all required input data is available. Args: start_datetime: First potential t0 time end_datetime: Last potential t0 time config_path: Path to data config file Returns: pandas.DatetimeIndex of the init-times available for required inputs """ start_datetime = pd.Timestamp(start_datetime) end_datetime = pd.Timestamp(end_datetime) # Open all the input data so we can check what of the potential data init times we have input # data for datapipes_dict = _get_datapipes_dict(config_path, production=False) # Pop out the config file config = datapipes_dict.pop("config") # We are going to abuse the `create_t0_and_loc_datapipes()` function to find the init-times in # potential_init_times which we have input data for. To do this, we will feed in some fake site # data which has the potential_init_times as timestamps. This is a bit hacky but works for now # Set up init-times we would like to make predictions for potential_init_times = pd.date_range(start_datetime, end_datetime, freq=f"{FREQ_MINS}min") # We buffer the potential init-times so that we don't lose any init-times from the # start and end. Again this is a hacky step history_duration = pd.Timedelta(config.input_data.pv.history_minutes, "min") forecast_duration = pd.Timedelta(config.input_data.pv.forecast_minutes, "min") buffered_potential_init_times = pd.date_range( start_datetime - history_duration, end_datetime + forecast_duration, freq=f"{FREQ_MINS}min" ) ds_fake_site = ( buffered_potential_init_times.to_frame().to_xarray().rename({"index": "time_utc"}) ) ds_fake_site = ds_fake_site.rename({0: "site_pv_power_mw"}) ds_fake_site = ds_fake_site.expand_dims("pv_system_id", axis=1) ds_fake_site = ds_fake_site.assign_coords( pv_system_id=[0], latitude=("pv_system_id", [0]), longitude=("pv_system_id", [0]), ) ds_fake_site = ds_fake_site.site_pv_power_mw.astype(float) * 1e-18 # Overwrite the site data which is already in the datapipes dict datapipes_dict["pv"] = IterableWrapper([ds_fake_site]) # Use create_t0_and_loc_datapipes to get datapipe of init-times location_pipe, t0_datapipe = create_t0_and_loc_datapipes( datapipes_dict, configuration=config, key_for_t0="pv", shuffle=False, ) # Create a full list of available init-times. Note that we need to loop over the t0s AND # locations to avoid the torch datapipes buffer overflow but we don't actually use the location available_init_times = [t0 for _, t0 in zip(location_pipe, t0_datapipe)] available_init_times = pd.to_datetime(available_init_times) logger.info( f"{len(available_init_times)} out of {len(potential_init_times)} " "requested init-times have required input data" ) return available_init_times def get_loctimes_datapipes(config_path): """Create location and init-time datapipes Args: config_path: Path to data config file Returns: tuple: A tuple of datapipes - Datapipe yielding locations - Datapipe yielding init-times """ # Set up ID location query object ds_sites = get_sites_ds(config_path) site_id_to_loc = SiteLocationLookup(ds_sites.longitude, ds_sites.latitude) # Filter the init-times to times we have all input data for available_target_times = get_available_t0_times( start_datetime, end_datetime, config_path, ) num_t0s = len(available_target_times) # Save the init-times which predictions are being made for. This is really helpful to check # whilst the backtest is running since it takes a long time. This lets you see what init-times # the backtest will end up producing available_target_times.to_frame().to_csv(f"{output_dir}/t0_times.csv") # Cycle the site locations location_pipe = IterableWrapper([[site_id_to_loc(site_id) for site_id in ALL_SITE_IDS]]).repeat( num_t0s ) # Shard and then unbatch the locations so that each worker will generate all samples for all # sites and for a single init-time location_pipe = location_pipe.sharding_filter() location_pipe = location_pipe.unbatch( unbatch_level=1 ) # might not need this part since the site datapipe is creating examples # Create times datapipe so each worker receives # len(ALL_SITE_IDS) copies of the same datetime for its batch t0_datapipe = IterableWrapper( [[t0 for site_id in ALL_SITE_IDS] for t0 in available_target_times] ) t0_datapipe = t0_datapipe.sharding_filter() t0_datapipe = t0_datapipe.unbatch( unbatch_level=1 ) # might not need this part since the site datapipe is creating examples t0_datapipe = t0_datapipe.set_length(num_t0s * len(ALL_SITE_IDS)) location_pipe = location_pipe.set_length(num_t0s * len(ALL_SITE_IDS)) return location_pipe, t0_datapipe class ModelPipe: """A class to conveniently make and process predictions from batches""" def __init__(self, model, ds_site: xr.Dataset): """A class to conveniently make and process predictions from batches Args: model: PVNet site level model ds_site:xarray dataset of pv site true values and capacities """ self.model = model self.ds_site = ds_site def predict_batch(self, batch: NumpyBatch) -> xr.Dataset: """Run the batch through the model and compile the predictions into an xarray DataArray Args: batch: A batch of samples with inputs for each site for the same init-time Returns: xarray.Dataset of all site and national forecasts for the batch """ # Unpack some variables from the batch id0 = batch[BatchKey.pv_t0_idx] t0 = batch[BatchKey.pv_time_utc].cpu().numpy().astype("datetime64[s]")[0, id0] n_valid_times = len(batch[BatchKey.pv_time_utc][0, id0 + 1 :]) model = self.model # Get valid times for this forecast valid_times = pd.to_datetime( [t0 + np.timedelta64((i + 1) * FREQ_MINS, "m") for i in range(n_valid_times)] ) # Get effective capacities for this forecast site_capacities = self.ds_site.nominal_capacity_wp.values # Get the solar elevations. We need to un-normalise these from the values in the batch elevation = batch[BatchKey.pv_solar_elevation] * ELEVATION_STD + ELEVATION_MEAN # We only need elevation mask for forecasted values, not history elevation = elevation[:, id0 + 1 :] # Make mask dataset for sundown da_sundown_mask = xr.DataArray( data=elevation < MIN_DAY_ELEVATION, dims=["pv_system_id", "target_datetime_utc"], coords=dict( pv_system_id=ALL_SITE_IDS, target_datetime_utc=valid_times, ), ) with torch.no_grad(): # Run batch through model to get 0-1 predictions for all sites device_batch = copy_batch_to_device(batch_to_tensor(batch), device) y_normed_site = model(device_batch).detach().cpu().numpy() da_normed_site = preds_to_dataarray(y_normed_site, model, valid_times, ALL_SITE_IDS) # Multiply normalised forecasts by capacities and clip negatives da_abs_site = da_normed_site.clip(0, None) * site_capacities[:, None, None] # Apply sundown mask da_abs_site = da_abs_site.where(~da_sundown_mask).fillna(0.0) da_abs_site = da_abs_site.expand_dims(dim="init_time_utc", axis=0).assign_coords( init_time_utc=np.array([t0], dtype="datetime64[ns]") ) return da_abs_site def get_datapipe(config_path: str) -> NumpyBatch: """Construct datapipe yielding batches of concurrent samples for all sites Args: config_path: Path to the data configuration file Returns: NumpyBatch: Concurrent batch of samples for each site """ # Construct location and init-time datapipes location_pipe, t0_datapipe = get_loctimes_datapipes(config_path) # Get the number of init-times # num_batches = len(t0_datapipe) num_batches = len(t0_datapipe) // len(ALL_SITE_IDS) # Construct sample datapipes data_pipeline = construct_sliced_data_pipeline( config_path, location_pipe, t0_datapipe, ) config = load_yaml_configuration(config_path) data_pipeline["pv"] = data_pipeline["pv"].pad_forward_pv( forecast_duration=np.timedelta64(config.input_data.pv.forecast_minutes, "m"), history_duration=np.timedelta64(config.input_data.pv.history_minutes, "m"), time_resolution_minutes=np.timedelta64(config.input_data.pv.time_resolution_minutes, "m"), ) data_pipeline = DictDatasetIterDataPipe( {k: v for k, v in data_pipeline.items() if k != "config"}, ).map(split_dataset_dict_dp) data_pipeline = data_pipeline.pvnet_site_convert_to_numpy_batch() # Batch so that each worker returns a batch of all locations for a single init-time # Also convert to tensor for model data_pipeline = ( data_pipeline.batch(len(ALL_SITE_IDS)) .map(stack_np_examples_into_batch) .map(batch_to_tensor) ) data_pipeline = data_pipeline.set_length(num_batches) return data_pipeline @hydra.main(config_path="../configs", config_name="config.yaml", version_base="1.2") def main(config: DictConfig): """Runs the backtest""" dataloader_kwargs = dict( shuffle=False, batch_size=None, sampler=None, batch_sampler=None, # Number of workers set in the config file num_workers=config.datamodule.num_workers, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, prefetch_factor=config.datamodule.prefetch_factor, persistent_workers=False, ) # Set up output dir os.makedirs(output_dir) # Create concurrent batch datapipe # Each batch includes a sample for each of the n sites for a single init-time batch_pipe = get_datapipe(config.datamodule.configuration) num_batches = len(batch_pipe) # Load the site data as an xarray object ds_site = get_sites_ds(config.datamodule.configuration) # Create a dataloader for the concurrent batches and use multiprocessing dataloader = DataLoader(batch_pipe, **dataloader_kwargs) # Load the PVNet model if model_chckpoint_dir: model, *_ = get_model_from_checkpoints([model_chckpoint_dir], val_best=True) elif hf_model_id: model = load_model_from_hf(hf_model_id, hf_revision, hf_token) else: raise ValueError("Provide a model checkpoint or a HuggingFace model") model = model.eval().to(device) # Create object to make predictions for each input batch model_pipe = ModelPipe(model, ds_site) # Loop through the batches pbar = tqdm(total=num_batches) for i, batch in zip(range(num_batches), dataloader): try: # Make predictions for the init-time ds_abs_all = model_pipe.predict_batch(batch) t0 = ds_abs_all.init_time_utc.values[0] # Save the predictions filename = f"{output_dir}/{t0}.nc" ds_abs_all.to_netcdf(filename) pbar.update() except Exception as e: print(f"Exception {e} at batch {i}") pass # Close down pbar.close() del dataloader if __name__ == "__main__": main()