PVNet_summation 1.0.0__py3-none-any.whl

pvnet_summation/__init__.py

@@ -0,0 +1 @@
+"""PVNet_summation"""

pvnet_summation/data/__init__.py

@@ -0,0 +1,2 @@
+"""Data module"""
+from .datamodule import PresavedDataModule, StreamedDataModule

pvnet_summation/data/datamodule.py

@@ -0,0 +1,213 @@
+"""Pytorch lightning datamodules for loading pre-saved samples and predictions."""
+
+from glob import glob
+from typing import TypeAlias
+
+import numpy as np
+import pandas as pd
+import torch
+from lightning.pytorch import LightningDataModule
+from ocf_data_sampler.load.gsp import open_gsp
+from ocf_data_sampler.numpy_sample.common_types import NumpyBatch, NumpySample
+from ocf_data_sampler.torch_datasets.datasets.pvnet_uk import PVNetUKConcurrentDataset
+from ocf_data_sampler.utils import minutes
+from torch.utils.data import DataLoader, Dataset, default_collate
+from typing_extensions import override
+
+SumNumpySample: TypeAlias = dict[str, np.ndarray | NumpyBatch]
+SumTensorBatch: TypeAlias = dict[str, torch.Tensor]
+
+
+class StreamedDataset(PVNetUKConcurrentDataset):
+    """A torch dataset for creating concurrent PVNet inputs and national targets."""
+
+    def __init__(
+        self,
+        config_filename: str,
+        start_time: str | None = None,
+        end_time: str | None = None,
+    ) -> None:
+        """A torch dataset for creating concurrent PVNet inputs and national targets.
+
+        Args:
+            config_filename: Path to the configuration file
+            start_time: Limit the init-times to be after this
+            end_time: Limit the init-times to be before this
+        """
+        super().__init__(config_filename, start_time, end_time, gsp_ids=None)
+
+        # Load and nornmalise the national GSP data to use as target values
+        national_gsp_data = (
+            open_gsp(
+                zarr_path=self.config.input_data.gsp.zarr_path, 
+                boundaries_version=self.config.input_data.gsp.boundaries_version
+            )
+            .sel(gsp_id=0)
+            .compute()
+        )
+        self.national_gsp_data = national_gsp_data / national_gsp_data.effective_capacity_mwp
+
+
+    def _get_sample(self, t0: pd.Timestamp) -> SumNumpySample:
+        """Generate a concurrent PVNet sample for given init-time.
+
+        Args:
+            t0: init-time for sample
+        """
+
+        pvnet_inputs: NumpySample = super()._get_sample(t0)
+
+        location_capacities = pvnet_inputs["gsp_effective_capacity_mwp"]
+
+        valid_times = pd.date_range(
+            t0+minutes(self.config.input_data.gsp.time_resolution_minutes), 
+            t0+minutes(self.config.input_data.gsp.interval_end_minutes),
+            freq=minutes(self.config.input_data.gsp.time_resolution_minutes)
+        )
+
+        total_outturns = self.national_gsp_data.sel(time_utc=valid_times).values
+        total_capacity = self.national_gsp_data.sel(time_utc=t0).effective_capacity_mwp.item()
+
+        relative_capacities = location_capacities / total_capacity
+
+        return {
+            # NumpyBatch object with batch size = num_locations
+            "pvnet_inputs": pvnet_inputs,
+            # Shape: [time]
+            "target": total_outturns,
+            # Shape: [time]
+            "valid_times": valid_times.values.astype(int),
+            # Shape: 
+            "last_outturn": self.national_gsp_data.sel(time_utc=t0).values,
+            # Shape: [num_locations]
+            "relative_capacity": relative_capacities,
+        }
+
+    @override
+    def __getitem__(self, idx: int) -> SumNumpySample:
+        return super().__getitem__(idx)
+
+    @override
+    def get_sample(self, t0: pd.Timestamp) -> SumNumpySample:
+        return super().get_sample(t0)
+
+
+class StreamedDataModule(LightningDataModule):
+    """Datamodule for training pvnet_summation."""
+
+    def __init__(
+        self,
+        configuration: str,
+        train_period: list[str | None] = [None, None],
+        val_period: list[str | None] = [None, None],
+        num_workers: int = 0,
+        prefetch_factor: int | None = None,
+        persistent_workers: bool = False,
+    ):
+        """Datamodule for creating concurrent PVNet inputs and national targets.
+
+        Args:
+            configuration: Path to ocf-data-sampler configuration file.
+            train_period: Date range filter for train dataloader.
+            val_period: Date range filter for val dataloader.
+            num_workers: Number of workers to use in multiprocess batch loading.
+            prefetch_factor: Number of data will be prefetched at the end of each worker process.
+            persistent_workers: If True, the data loader will not shut down the worker processes 
+                after a dataset has been consumed once. This allows to maintain the workers Dataset 
+                instances alive.
+        """
+        super().__init__()
+        self.configuration = configuration
+        self.train_period = train_period
+        self.val_period = val_period
+
+        self._dataloader_kwargs = dict(
+            batch_size=None,
+            batch_sampler=None,
+            num_workers=num_workers,
+            collate_fn=None,
+            pin_memory=False,
+            drop_last=False,
+            timeout=0,
+            worker_init_fn=None,
+            prefetch_factor=prefetch_factor,
+            persistent_workers=persistent_workers,
+        )
+
+    def train_dataloader(self, shuffle: bool = False) -> DataLoader:
+        """Construct train dataloader"""
+        dataset = StreamedDataset(self.configuration, *self.train_period)
+        return DataLoader(dataset, shuffle=shuffle, **self._dataloader_kwargs)
+
+    def val_dataloader(self, shuffle: bool = False) -> DataLoader:
+        """Construct val dataloader"""
+        dataset = StreamedDataset(self.configuration, *self.val_period)
+        return DataLoader(dataset, shuffle=shuffle, **self._dataloader_kwargs)
+
+
+class PresavedDataset(Dataset):
+    """Dataset for loading pre-saved PVNet predictions from disk"""
+
+    def __init__(self, sample_dir: str):
+        """"Dataset for loading pre-saved PVNet predictions from disk.
+        
+        Args:
+            sample_dir: The directory containing the saved samples
+        """
+        self.sample_filepaths = sorted(glob(f"{sample_dir}/*.pt"))
+
+    def __len__(self) -> int:
+        return len(self.sample_filepaths)
+
+    def __getitem__(self, idx: int) -> dict:
+        return torch.load(self.sample_filepaths[idx], weights_only=True)
+
+
+class PresavedDataModule(LightningDataModule):
+    """Datamodule for loading pre-saved PVNet predictions."""
+
+    def __init__(
+        self,
+        sample_dir: str,
+        batch_size: int = 16,
+        num_workers: int = 0,
+        prefetch_factor: int | None = None,
+        persistent_workers: bool = False,
+    ):
+        """Datamodule for loading pre-saved PVNet predictions.
+
+        Args:
+            sample_dir: Path to the directory of pre-saved samples.
+            batch_size: Batch size.
+            num_workers: Number of workers to use in multiprocess batch loading.
+            prefetch_factor: Number of data will be prefetched at the end of each worker process.
+            persistent_workers: If True, the data loader will not shut down the worker processes 
+                after a dataset has been consumed once. This allows to maintain the workers Dataset 
+                instances alive.
+        """
+        super().__init__()
+        self.sample_dir = sample_dir
+
+        self._dataloader_kwargs = dict(
+            batch_size=batch_size,
+            sampler=None,
+            batch_sampler=None,
+            num_workers=num_workers,
+            collate_fn=None if batch_size is None else default_collate,
+            pin_memory=False,
+            drop_last=False,
+            timeout=0,
+            worker_init_fn=None,
+            prefetch_factor=prefetch_factor,
+            persistent_workers=persistent_workers,
+        )
+
+    def train_dataloader(self, shuffle: bool = True) -> DataLoader:
+        """Construct train dataloader"""
+        dataset = PresavedDataset(f"{self.sample_dir}/train")
+        return DataLoader(dataset, shuffle=shuffle, **self._dataloader_kwargs)
+
+    def val_dataloader(self, shuffle: bool = False) -> DataLoader:
+        """Construct val dataloader"""
+        dataset = PresavedDataset(f"{self.sample_dir}/val")
+        return DataLoader(dataset, shuffle=shuffle, **self._dataloader_kwargs)

pvnet_summation/load_model.py

@@ -0,0 +1,70 @@
+"""Load a model from its checkpoint directory"""
+
+import glob
+import os
+
+import hydra
+import torch
+import yaml
+
+from pvnet_summation.utils import (
+    DATAMODULE_CONFIG_NAME,
+    FULL_CONFIG_NAME,
+    MODEL_CONFIG_NAME,
+)
+
+
+def get_model_from_checkpoints(
+    checkpoint_dir_path: str,
+    val_best: bool = True,
+) -> tuple[torch.nn.Module, dict, str | None, str | None]:
+    """Load a model from its checkpoint directory
+
+    Returns:
+        tuple:
+            model: nn.Module of pretrained model.
+            model_config: path to model config used to train the model.
+            datamodule_config: path to datamodule used to create samples e.g train/test split info.
+            experiment_configs: path to the full experimental config.
+
+    """
+
+    # Load lightning training module
+    with open(f"{checkpoint_dir_path}/{MODEL_CONFIG_NAME}") as cfg:
+        model_config = yaml.load(cfg, Loader=yaml.FullLoader)
+
+    lightning_module = hydra.utils.instantiate(model_config)
+
+    if val_best:
+        # Only one epoch (best) saved per model
+        files = glob.glob(f"{checkpoint_dir_path}/epoch*.ckpt")
+        if len(files) != 1:
+            raise ValueError(
+                f"Found {len(files)} checkpoints @ {checkpoint_dir_path}/epoch*.ckpt. Expected one."
+            )
+        
+        checkpoint = torch.load(files[0], map_location="cpu", weights_only=True)
+    else:
+        checkpoint = torch.load(
+            f"{checkpoint_dir_path}/last.ckpt", 
+            map_location="cpu", 
+            weights_only=True,
+        )
+
+    lightning_module.load_state_dict(state_dict=checkpoint["state_dict"])
+
+    # Extract the model from the lightning module
+    model = lightning_module.model
+    model_config = model_config["model"]
+
+    # Check for datamodule config
+    # This only exists if the model was trained with presaved samples
+    datamodule_config = f"{checkpoint_dir_path}/{DATAMODULE_CONFIG_NAME}"
+    datamodule_config = datamodule_config if os.path.isfile(datamodule_config) else None
+
+    # Check for experiment config
+    # For backwards compatibility - this might not always exist
+    experiment_config = f"{checkpoint_dir_path}/{FULL_CONFIG_NAME}"
+    experiment_config = experiment_config if os.path.isfile(experiment_config) else None
+
+    return model, model_config, datamodule_config, experiment_config

pvnet_summation/models/__init__.py

@@ -0,0 +1,3 @@
+"""Models for PVNet summation"""
+from .base_model import BaseModel
+from .dense_model import DenseModel

pvnet_summation/models/base_model.py

@@ -0,0 +1,345 @@
+"""Base model for all PVNet submodels"""
+import logging
+import os
+import shutil
+import time
+from importlib.metadata import version
+from pathlib import Path
+
+import hydra
+import torch
+import yaml
+from huggingface_hub import ModelCard, ModelCardData, snapshot_download
+from huggingface_hub.hf_api import HfApi
+from safetensors.torch import load_file, save_file
+
+from pvnet_summation.data.datamodule import SumTensorBatch
+from pvnet_summation.utils import (
+    DATAMODULE_CONFIG_NAME,
+    FULL_CONFIG_NAME,
+    MODEL_CARD_NAME,
+    MODEL_CONFIG_NAME,
+    PYTORCH_WEIGHTS_NAME,
+)
+
+
+def santize_datamodule(config: dict) -> dict:
+    """Create new datamodule config which only keeps the details required for inference"""
+    return {"pvnet_model": config["pvnet_model"]}
+
+
+def download_from_hf(
+    repo_id: str,
+    filename: str | list[str],
+    revision: str,
+    cache_dir: str | None,
+    force_download: bool,
+    max_retries: int = 5,
+    wait_time: int = 10,
+) -> str | list[str]:
+    """Tries to download one or more files from HuggingFace up to max_retries times.
+
+    Args:
+        repo_id: HuggingFace repo ID
+        filename: Name of the file(s) to download
+        revision: Specific model revision
+        cache_dir: Cache directory
+        force_download: Whether to force a new download
+        max_retries: Maximum number of retry attempts
+        wait_time: Wait time (in seconds) before retrying
+
+    Returns:
+        The local file path of the downloaded file(s)
+    """
+    for attempt in range(1, max_retries + 1):
+        try:
+            save_dir = snapshot_download(
+                repo_id=repo_id,
+                allow_patterns=filename,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+            )
+
+            if isinstance(filename, list):
+                return [f"{save_dir}/{f}" for f in filename]
+            else:
+                return f"{save_dir}/{filename}"
+        
+        except Exception as e:
+            if attempt == max_retries:
+                raise Exception(
+                    f"Failed to download {filename} from {repo_id} after {max_retries} attempts."
+                ) from e
+            logging.warning(
+                (
+                    f"Attempt {attempt}/{max_retries} failed to download {filename} "
+                    f"from {repo_id}. Retrying in {wait_time} seconds..."
+                )
+            )
+            time.sleep(wait_time)
+
+
+class HuggingfaceMixin:
+    """Mixin for saving and loading model to and from huggingface"""
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str,
+        revision: str,
+        cache_dir: str | None = None,
+        force_download: bool = False,
+        strict: bool = True,
+    ) -> "BaseModel":
+        """Load Pytorch pretrained weights and return the loaded model."""
+
+        if os.path.isdir(model_id):
+            print("Loading model from local directory")
+            model_file = f"{model_id}/{PYTORCH_WEIGHTS_NAME}"
+            config_file = f"{model_id}/{MODEL_CONFIG_NAME}"
+        else:
+            print("Loading model from huggingface repo")
+
+            model_file, config_file = download_from_hf(
+                repo_id=model_id,
+                filename=[PYTORCH_WEIGHTS_NAME, MODEL_CONFIG_NAME],
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                max_retries=5,
+                wait_time=10,
+            )
+
+        with open(config_file, "r") as f:
+            model = hydra.utils.instantiate(yaml.safe_load(f))
+
+        state_dict = load_file(model_file)
+        model.load_state_dict(state_dict, strict=strict)  # type: ignore
+        model.eval()  # type: ignore
+
+        return model
+    
+    @classmethod
+    def get_datamodule_config(
+        cls,
+        model_id: str,
+        revision: str,
+        cache_dir: str | None = None,
+        force_download: bool = False,
+    ) -> str:
+        """Load data config file."""
+        if os.path.isdir(model_id):
+            print("Loading datamodule config from local directory")
+            datamodule_config_file = os.path.join(model_id, DATAMODULE_CONFIG_NAME)
+        else:
+            print("Loading datamodule config from huggingface repo")
+            datamodule_config_file = download_from_hf(
+                repo_id=model_id,
+                filename=DATAMODULE_CONFIG_NAME,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                max_retries=5,
+                wait_time=10,
+            )
+
+        return datamodule_config_file
+
+    def _save_model_weights(self, save_directory: str) -> None:
+        """Save weights from a Pytorch model to a local directory."""
+        save_file(self.state_dict(), f"{save_directory}/{PYTORCH_WEIGHTS_NAME}")
+
+    def save_pretrained(
+        self,
+        save_directory: str,
+        model_config: dict,
+        wandb_repo: str,
+        wandb_id: str,
+        card_template_path: str,
+        datamodule_config_path,
+        experiment_config_path: str | None = None,
+        hf_repo_id: str | None = None,
+        push_to_hub: bool = False,
+    ) -> None:
+        """Save weights in local directory or upload to huggingface hub.
+
+        Args:
+            save_directory:
+                Path to directory in which the model weights and configuration will be saved.
+            model_config (`dict`):
+                Model configuration specified as a key/value dictionary.
+            wandb_repo: Identifier of the repo on wandb.
+            wandb_id: Identifier of the model on wandb.
+            datamodule_config_path:
+                The path to the datamodule config.
+            card_template_path: Path to the HuggingFace model card template. Defaults to card in
+                PVNet library if set to None.
+            experiment_config_path:
+                The path to the full experimental config.
+            hf_repo_id:
+                ID of your repository on the Hub. Used only if `push_to_hub=True`. Will default to
+                the folder name if not provided.
+            push_to_hub (`bool`, *optional*, defaults to `False`):
+                Whether or not to push your model to the HuggingFace Hub after saving it.
+        """
+
+        save_directory = Path(save_directory)
+        save_directory.mkdir(parents=True, exist_ok=True)
+
+        # Save model weights/files
+        self._save_model_weights(save_directory)
+
+        # Save the model config
+        if isinstance(model_config, dict):
+            with open(save_directory / MODEL_CONFIG_NAME, "w") as outfile:
+                yaml.dump(model_config, outfile, sort_keys=False, default_flow_style=False)
+
+        # Sanitize and save the datamodule config
+        with open(datamodule_config_path) as cfg:
+            datamodule_config = yaml.load(cfg, Loader=yaml.FullLoader)
+
+        datamodule_config = santize_datamodule(datamodule_config)
+
+        with open(save_directory / DATAMODULE_CONFIG_NAME, "w") as outfile:
+            yaml.dump(datamodule_config, outfile, sort_keys=False, default_flow_style=False)
+        
+        # Save the full experimental config
+        if experiment_config_path is not None:
+            shutil.copyfile(experiment_config_path, save_directory / FULL_CONFIG_NAME)
+
+        card = self.create_hugging_face_model_card(card_template_path, wandb_repo, wandb_id)
+
+        (save_directory / MODEL_CARD_NAME).write_text(str(card))
+
+        if push_to_hub:
+            api = HfApi()
+
+            api.upload_folder(
+                repo_id=hf_repo_id,
+                folder_path=save_directory,
+                repo_type="model",
+                commit_message=f"Upload model - {wandb_id}",
+            )
+
+            # Print the most recent commit hash
+            c = api.list_repo_commits(repo_id=hf_repo_id, repo_type="model")[0]
+
+            message = (
+                f"The latest commit is now: \n"
+                f"    date: {c.created_at} \n"
+                f"    commit hash: {c.commit_id}\n"
+                f"    by: {c.authors}\n"
+                f"    title: {c.title}\n"
+            )
+
+            print(message)
+
+    @staticmethod
+    def create_hugging_face_model_card(
+        card_template_path: str,
+        wandb_repo: str,
+        wandb_id: str,
+    ) -> ModelCard:
+        """
+        Creates Hugging Face model card
+
+        Args:
+            card_template_path: Path to the HuggingFace model card template
+            wandb_repo: Identifier of the repo on wandb.
+            wandb_id: Identifier of the model on wandb.
+
+        Returns:
+            card: ModelCard - Hugging Face model card object
+        """
+
+        # Creating and saving model card.
+        card_data = ModelCardData(language="en", license="mit", library_name="pytorch")
+
+        link = f"https://wandb.ai/{wandb_repo}/runs/{wandb_id}"
+        wandb_link = f" - [{link}]({link})\n"
+
+        # Find package versions for OCF packages
+        packages_to_display = ["pvnet_summation", "ocf-data-sampler"]
+        packages_and_versions = {package: version(package) for package in packages_to_display}
+
+
+        package_versions_markdown = ""
+        for package, v in packages_and_versions.items():
+            package_versions_markdown += f" - {package}=={v}\n"
+
+        return ModelCard.from_template(
+            card_data,
+            template_path=card_template_path,
+            wandb_link=wandb_link,
+            package_versions=package_versions_markdown,
+        )
+
+
+class BaseModel(torch.nn.Module, HuggingfaceMixin):
+    """Abstract base class for PVNet-summation submodels"""
+
+    def __init__(
+        self,
+        output_quantiles: list[float] | None,
+        num_input_locations: int,
+        input_quantiles: list[float] | None,
+        history_minutes: int,
+        forecast_minutes: int,
+        interval_minutes: int,
+    ):
+        """Abtstract base class for PVNet-summation submodels.
+
+        """
+        super().__init__()
+
+        self.output_quantiles = output_quantiles
+
+        self.num_input_locations = num_input_locations
+        self.input_quantiles = input_quantiles
+
+        self.history_minutes = history_minutes
+        self.forecast_minutes = forecast_minutes
+        self.interval_minutes = interval_minutes
+
+        # Number of timestemps for 30 minutely data
+        self.history_len = history_minutes // interval_minutes
+        self.forecast_len = (forecast_minutes) // interval_minutes
+
+        # Store whether the model should use quantile regression or simply predict the mean
+        self.use_quantile_regression = self.output_quantiles is not None
+
+        # Store the number of ouput features that the model should predict for
+        if self.use_quantile_regression:
+            self.num_output_features = self.forecast_len * len(self.output_quantiles)
+        else:
+            self.num_output_features = self.forecast_len
+
+        # Store the expected input shape
+        if input_quantiles is None:
+            self.input_shape = (self.num_input_locations, self.forecast_len)
+        else:
+            self.input_shape = (self.num_input_locations, self.forecast_len, len(input_quantiles))
+
+    def _quantiles_to_prediction(self, y_quantiles: torch.Tensor) -> torch.Tensor:
+        """Convert network prediction into a point prediction.
+
+        Args:
+            y_quantiles: Quantile prediction of network
+
+        Returns:
+            torch.Tensor: Point prediction
+        """
+        # y_quantiles Shape: [batch_size, seq_length, num_quantiles]
+        idx = self.output_quantiles.index(0.5)
+        return y_quantiles[..., idx]
+
+    def sum_of_locations(self, x: SumTensorBatch) -> torch.Tensor:
+        """Compute the sum of the location-level predictions"""
+        if self.input_quantiles is None:
+            y_hat = x["pvnet_outputs"]
+        else:
+            idx = self.input_quantiles.index(0.5)
+            y_hat = x["pvnet_outputs"][..., idx]
+
+        return (y_hat * x["relative_capacity"].unsqueeze(-1)).sum(dim=1)