PVNet 5.2.3__py3-none-any.whl → 5.3.0__py3-none-any.whl

pvnet/datamodule.py

@@ -1,4 +1,4 @@
-""" Data module for pytorch lightning """
+"""Data module for pytorch lightning"""
 
 import os
 
@@ -6,10 +6,9 @@ import numpy as np
 from lightning.pytorch import LightningDataModule
 from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
 from ocf_data_sampler.numpy_sample.common_types import NumpySample, TensorBatch
-from ocf_data_sampler.torch_datasets.datasets.pvnet_uk import PVNetUKRegionalDataset
-from ocf_data_sampler.torch_datasets.datasets.site import SitesDataset
+from ocf_data_sampler.torch_datasets.pvnet_dataset import PVNetDataset
 from ocf_data_sampler.torch_datasets.utils.torch_batch_utils import batch_to_tensor
-from torch.utils.data import DataLoader, Dataset, Subset
+from torch.utils.data import DataLoader, Subset
 
 
 def collate_fn(samples: list[NumpySample]) -> TensorBatch:
@@ -17,7 +16,7 @@ def collate_fn(samples: list[NumpySample]) -> TensorBatch:
     return batch_to_tensor(stack_np_samples_into_batch(samples))
 
 
-class BaseDataModule(LightningDataModule):
+class PVNetDataModule(LightningDataModule):
     """Base Datamodule which streams samples using a sampler from ocf-data-sampler."""
 
     def __init__(
@@ -40,10 +39,10 @@ class BaseDataModule(LightningDataModule):
             batch_size: Batch size.
             num_workers: Number of workers to use in multiprocess batch loading.
             prefetch_factor: Number of batches loaded in advance by each worker.
-            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 
+            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.
-            pin_memory: If True, the data loader will copy Tensors into device/CUDA pinned memory 
+            pin_memory: If True, the data loader will copy Tensors into device/CUDA pinned memory
                 before returning them.
             train_period: Date range filter for train dataloader.
             val_period: Date range filter for val dataloader.
@@ -70,7 +69,7 @@ class BaseDataModule(LightningDataModule):
             worker_init_fn=None,
             prefetch_factor=prefetch_factor,
             persistent_workers=persistent_workers,
-            multiprocessing_context="spawn" if num_workers>0 else None,
+            multiprocessing_context="spawn" if num_workers > 0 else None,
         )
 
     def setup(self, stage: str | None = None):
@@ -79,16 +78,15 @@ class BaseDataModule(LightningDataModule):
         # This logic runs only once at the start of training, therefore the val dataset is only
         # shuffled once
         if stage == "fit":
-
             # Prepare the train dataset
             self.train_dataset = self._get_dataset(*self.train_period)
 
-            # Prepare and pre-shuffle the val dataset and set seed for reproducibility
+            # Prepare and pre-shuffle the val dataset and set seed for reproducibility
             val_dataset = self._get_dataset(*self.val_period)
 
             shuffled_indices = np.random.default_rng(seed=self.seed).permutation(len(val_dataset))
             self.val_dataset = Subset(val_dataset, shuffled_indices)
-        
+
             if self.dataset_pickle_dir is not None:
                 os.makedirs(self.dataset_pickle_dir, exist_ok=True)
                 train_dataset_path = f"{self.dataset_pickle_dir}/train_dataset.pkl"
@@ -116,8 +114,8 @@ class BaseDataModule(LightningDataModule):
                 if os.path.exists(filepath):
                     os.remove(filepath)
 
-    def _get_dataset(self, start_time: str | None, end_time: str | None) -> Dataset:
-        raise NotImplementedError
+    def _get_dataset(self, start_time: str | None, end_time: str | None) -> PVNetDataset:
+        return PVNetDataset(self.configuration, start_time=start_time, end_time=end_time)
 
     def train_dataloader(self) -> DataLoader:
         """Construct train dataloader"""
@@ -126,17 +124,3 @@ class BaseDataModule(LightningDataModule):
     def val_dataloader(self) -> DataLoader:
         """Construct val dataloader"""
         return DataLoader(self.val_dataset, shuffle=False, **self._common_dataloader_kwargs)
-
-
-class UKRegionalDataModule(BaseDataModule):
-    """Datamodule for streaming UK regional samples."""
-
-    def _get_dataset(self, start_time: str | None, end_time: str | None) -> PVNetUKRegionalDataset:
-        return PVNetUKRegionalDataset(self.configuration, start_time=start_time, end_time=end_time)
-
-
-class SitesDataModule(BaseDataModule):
-    """Datamodule for streaming site samples."""
-
-    def _get_dataset(self, start_time: str | None, end_time: str | None) -> SitesDataset:
-        return SitesDataset(self.configuration, start_time=start_time, end_time=end_time)

pvnet/models/base_model.py

@@ -1,4 +1,5 @@
 """Base model for all PVNet submodels"""
+
 import logging
 import os
 import shutil
@@ -32,7 +33,7 @@ def fill_config_paths_with_placeholder(config: dict, placeholder: str = "PLACEHO
     """
     input_config = config["input_data"]
 
-    for source in ["gsp", "satellite"]:
+    for source in ["generation", "satellite"]:
         if source in input_config:
             # If not empty - i.e. if used
             if input_config[source]["zarr_path"] != "":
@@ -78,8 +79,8 @@ def minimize_config_for_model(config: dict, model: "BaseModel") -> dict:
 
                     # Replace the interval_end_minutes minutes
                     nwp_config["interval_end_minutes"] = (
-                        nwp_config["interval_start_minutes"] +
-                        (model.nwp_encoders_dict[nwp_source].sequence_length - 1)
+                        nwp_config["interval_start_minutes"]
+                        + (model.nwp_encoders_dict[nwp_source].sequence_length - 1)
                         * nwp_config["time_resolution_minutes"]
                     )
 
@@ -96,20 +97,19 @@ def minimize_config_for_model(config: dict, model: "BaseModel") -> dict:
 
             # Replace the interval_end_minutes minutes
             sat_config["interval_end_minutes"] = (
-                sat_config["interval_start_minutes"] +
-                (model.sat_encoder.sequence_length - 1)
-                * sat_config["time_resolution_minutes"]
+                sat_config["interval_start_minutes"]
+                + (model.sat_encoder.sequence_length - 1) * sat_config["time_resolution_minutes"]
             )
 
     if "pv" in input_config:
         if not model.include_pv:
             del input_config["pv"]
 
-    if "gsp" in input_config:
-        gsp_config = input_config["gsp"]
+    if "generation" in input_config:
+        generation_config = input_config["generation"]
 
         # Replace the forecast minutes
-        gsp_config["interval_end_minutes"] = model.forecast_minutes
+        generation_config["interval_end_minutes"] = model.forecast_minutes
 
     if "solar_position" in input_config:
         solar_config = input_config["solar_position"]
@@ -138,9 +138,9 @@ def download_from_hf(
         force_download: Whether to force a new download
         max_retries: Maximum number of retry attempts
         wait_time: Wait time (in seconds) before retrying
-        token: 
+        token:
             HF authentication token. If True, the token is read from the HuggingFace config folder.
-            If a string, it is used as the authentication token. 
+            If a string, it is used as the authentication token.
 
     Returns:
         The local file path of the downloaded file(s)
@@ -160,7 +160,7 @@ def download_from_hf(
                 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(
@@ -205,7 +205,7 @@ class HuggingfaceMixin:
                 force_download=force_download,
                 max_retries=5,
                 wait_time=10,
-                token=token
+                token=token,
             )
 
         with open(config_file, "r") as f:
@@ -240,7 +240,7 @@ class HuggingfaceMixin:
                 force_download=force_download,
                 max_retries=5,
                 wait_time=10,
-                token=token
+                token=token,
             )
 
         return data_config_file
@@ -301,7 +301,7 @@ class HuggingfaceMixin:
         # Save cleaned version of input data configuration file
         with open(data_config_path) as cfg:
             config = yaml.load(cfg, Loader=yaml.FullLoader)
-        
+
         config = fill_config_paths_with_placeholder(config)
         config = minimize_config_for_model(config, self)
 
@@ -311,7 +311,7 @@ class HuggingfaceMixin:
         # Save the datamodule config
         if datamodule_config_path is not None:
             shutil.copyfile(datamodule_config_path, save_directory / DATAMODULE_CONFIG_NAME)
-        
+
         # Save the full experimental config
         if experiment_config_path is not None:
             shutil.copyfile(experiment_config_path, save_directory / FULL_CONFIG_NAME)
@@ -378,7 +378,6 @@ class HuggingfaceMixin:
         packages_to_display = ["pvnet", "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"
@@ -399,23 +398,19 @@ class BaseModel(torch.nn.Module, HuggingfaceMixin):
         history_minutes: int,
         forecast_minutes: int,
         output_quantiles: list[float] | None = None,
-        target_key: str = "gsp",
         interval_minutes: int = 30,
     ):
         """Abtstract base class for PVNet submodels.
 
         Args:
-            history_minutes (int): Length of the GSP history period in minutes
-            forecast_minutes (int): Length of the GSP forecast period in minutes
+            history_minutes (int): Length of the generation history period in minutes
+            forecast_minutes (int): Length of the generation forecast period in minutes
             output_quantiles: A list of float (0.0, 1.0) quantiles to predict values for. If set to
                 None the output is a single value.
-            target_key: The key of the target variable in the batch
             interval_minutes: The interval in minutes between each timestep in the data
         """
         super().__init__()
 
-        self._target_key = target_key
-
         self.history_minutes = history_minutes
         self.forecast_minutes = forecast_minutes
         self.output_quantiles = output_quantiles

pvnet/models/ensemble.py

@@ -26,7 +26,6 @@ class Ensemble(BaseModel):
         output_quantiles = []
         history_minutes = []
         forecast_minutes = []
-        target_key = []
         interval_minutes = []
 
         # Get some model properties from each model
@@ -34,7 +33,6 @@ class Ensemble(BaseModel):
             output_quantiles.append(model.output_quantiles)
             history_minutes.append(model.history_minutes)
             forecast_minutes.append(model.forecast_minutes)
-            target_key.append(model._target_key)
             interval_minutes.append(model.interval_minutes)
 
         # Check these properties are all the same
@@ -42,7 +40,6 @@ class Ensemble(BaseModel):
             output_quantiles,
             history_minutes,
             forecast_minutes,
-            target_key,
             interval_minutes,
         ]:
             assert all([p == param_list[0] for p in param_list]), param_list
@@ -51,7 +48,6 @@ class Ensemble(BaseModel):
             history_minutes=history_minutes[0],
             forecast_minutes=forecast_minutes[0],
             output_quantiles=output_quantiles[0],
-            target_key=target_key[0],
             interval_minutes=interval_minutes[0],
         )
 

pvnet/models/late_fusion/late_fusion.py

@@ -28,8 +28,8 @@ class LateFusionModel(BaseModel):
     - NWP, if included, is put through a similar encoder.
     - PV site-level data, if included, is put through an encoder which transforms it from 2D, with
         time and system-ID dimensions, to become a 1D feature vector.
-    - The satellite features*, NWP features*, PV site-level features*, GSP ID embedding*, and sun
-        paramters* are concatenated into a 1D feature vector and passed through another neural
+    - The satellite features*, NWP features*, PV site-level features*, location ID embedding*, and
+        sun paramters* are concatenated into a 1D feature vector and passed through another neural
         network to combine them and produce a forecast.
 
     * if included
@@ -43,8 +43,7 @@ class LateFusionModel(BaseModel):
         sat_encoder: AbstractNWPSatelliteEncoder | None = None,
         pv_encoder: AbstractSitesEncoder | None = None,
         add_image_embedding_channel: bool = False,
-        include_gsp_yield_history: bool = True,
-        include_site_yield_history: bool = False,
+        include_generation_history: bool = False,
         include_sun: bool = True,
         include_time: bool = False,
         location_id_mapping: dict[Any, int] | None = None,
@@ -56,7 +55,6 @@ class LateFusionModel(BaseModel):
         nwp_forecast_minutes: DictConfig | None = None,
         nwp_history_minutes: DictConfig | None = None,
         pv_history_minutes: int | None = None,
-        target_key: str = "gsp",
         interval_minutes: int = 30,
         nwp_interval_minutes: DictConfig | None = None,
         pv_interval_minutes: int = 5,
@@ -83,14 +81,13 @@ class LateFusionModel(BaseModel):
             pv_encoder: A partially instantiated pytorch Module class used to encode the site-level
                 PV data from 2D into a 1D feature vector.
             add_image_embedding_channel: Add a channel to the NWP and satellite data with the
-                embedding of the GSP ID.
-            include_gsp_yield_history: Include GSP yield data.
-            include_site_yield_history: Include Site yield data.
+                embedding of the location ID.
+            include_generation_history: Include generation yield data.
             include_sun: Include sun azimuth and altitude data.
             include_time: Include sine and cosine of dates and times.
             location_id_mapping: A dictionary mapping the location ID to an integer. ID embedding is
                 not used if this is not provided.
-            embedding_dim: Number of embedding dimensions to use for GSP ID.
+            embedding_dim: Number of embedding dimensions to use for location ID.
             forecast_minutes: The amount of minutes that should be forecasted.
             history_minutes: The default amount of historical minutes that are used.
             sat_history_minutes: Length of recent observations used for satellite inputs. Defaults
@@ -103,7 +100,6 @@ class LateFusionModel(BaseModel):
                 `history_minutes` if not provided.
             pv_history_minutes: Length of recent site-level PV data used as
                 input. Defaults to `history_minutes` if not provided.
-            target_key: The key of the target variable in the batch.
             interval_minutes: The interval between each sample of the target data
             nwp_interval_minutes: Dictionary of the intervals between each sample of the NWP
                 data for each source
@@ -114,12 +110,10 @@ class LateFusionModel(BaseModel):
             history_minutes=history_minutes,
             forecast_minutes=forecast_minutes,
             output_quantiles=output_quantiles,
-            target_key=target_key,
             interval_minutes=interval_minutes,
         )
 
-        self.include_gsp_yield_history = include_gsp_yield_history
-        self.include_site_yield_history = include_site_yield_history
+        self.include_generation_history = include_generation_history
         self.include_sat = sat_encoder is not None
         self.include_nwp = nwp_encoders_dict is not None and len(nwp_encoders_dict) != 0
         self.include_pv = pv_encoder is not None
@@ -133,8 +127,7 @@ class LateFusionModel(BaseModel):
 
         if self.location_id_mapping is None:
             logger.warning(
-                "location_id_mapping` is not provided, defaulting to outdated GSP mapping"
-                "(0 to 317)"
+                "location_id_mapping` is not provided, defaulting to outdated GSP mapping(0 to 317)"
             )
 
             # Note 318 is the 2024 UK GSP count, so this is a temporary fix
@@ -223,8 +216,7 @@ class LateFusionModel(BaseModel):
 
             self.pv_encoder = pv_encoder(
                 sequence_length=pv_history_minutes // pv_interval_minutes + 1,
-                target_key_to_use=self._target_key,
-                input_key_to_use="site",
+                key_to_use="generation",
             )
 
             # Update num features
@@ -238,8 +230,7 @@ class LateFusionModel(BaseModel):
 
         if self.include_sun:
             self.sun_fc1 = nn.Linear(
-                in_features=2
-                * (self.forecast_len + self.history_len + 1),
+                in_features=2 * (self.forecast_len + self.history_len + 1),
                 out_features=16,
             )
 
@@ -248,19 +239,14 @@ class LateFusionModel(BaseModel):
 
         if self.include_time:
             self.time_fc1 = nn.Linear(
-                in_features=4
-                * (self.forecast_len + self.history_len + 1),
+                in_features=4 * (self.forecast_len + self.history_len + 1),
                 out_features=32,
             )
 
             # Update num features
             fusion_input_features += 32
 
-        if include_gsp_yield_history:
-            # Update num features
-            fusion_input_features += self.history_len
-
-        if include_site_yield_history:
+        if include_generation_history:
             # Update num features
             fusion_input_features += self.history_len + 1
 
@@ -269,15 +255,14 @@ class LateFusionModel(BaseModel):
             out_features=self.num_output_features,
         )
 
-
     def forward(self, x: TensorBatch) -> torch.Tensor:
         """Run model forward"""
 
         if self.use_id_embedding:
-            # eg: x['gsp_id'] = [1] with location_id_mapping = {1:0}, would give [0]
+            # eg: x['location_id'] = [1] with location_id_mapping = {1:0}, would give [0]
             id = torch.tensor(
-                [self.location_id_mapping[i.item()] for i in x[f"{self._target_key}_id"]],
-                device=x[f"{self._target_key}_id"].device,
+                [self.location_id_mapping[i.item()] for i in x["location_id"]],
+                device=x["location_id"].device,
                 dtype=torch.int64,
             )
 
@@ -308,32 +293,20 @@ class LateFusionModel(BaseModel):
                 nwp_out = self.nwp_encoders_dict[nwp_source](nwp_data)
                 modes[f"nwp/{nwp_source}"] = nwp_out
 
-        # *********************** Site Data *************************************
-        # Add site-level yield history
-        if self.include_site_yield_history:
-            site_history = x["site"][:, : self.history_len + 1].float()
-            site_history = site_history.reshape(site_history.shape[0], -1)
-            modes["site"] = site_history
+        # *********************** Generation Data *************************************
+        # Add generation yield history
+        if self.include_generation_history:
+            generation_history = x["generation"][:, : self.history_len + 1].float()
+            generation_history = generation_history.reshape(generation_history.shape[0], -1)
+            modes["generation"] = generation_history
 
-        # Add site-level yield history through PV encoder
+        # Add location-level yield history through PV encoder
         if self.include_pv:
-            if self._target_key != "site":
-                modes["site"] = self.pv_encoder(x)
-            else:
-                # Target is PV, so only take the history
-                # Copy batch
-                x_tmp = x.copy()
-                x_tmp["site"] = x_tmp["site"][:, : self.history_len + 1]
-                modes["site"] = self.pv_encoder(x_tmp)
-
-        # *********************** GSP Data ************************************
-        # Add gsp yield history
-        if self.include_gsp_yield_history:
-            gsp_history = x["gsp"][:, : self.history_len].float()
-            gsp_history = gsp_history.reshape(gsp_history.shape[0], -1)
-            modes["gsp"] = gsp_history
-
-        # ********************** Embedding of GSP/Site ID ********************
+            x_tmp = x.copy()
+            x_tmp["generation"] = x_tmp["generation"][:, : self.history_len + 1]
+            modes["generation"] = self.pv_encoder(x_tmp)
+
+        # ********************** Embedding of location ID ********************
         if self.use_id_embedding:
             modes["id"] = self.embed(id)
 
@@ -341,7 +314,7 @@ class LateFusionModel(BaseModel):
             sun = torch.cat((x["solar_azimuth"], x["solar_elevation"]), dim=1).float()
             sun = self.sun_fc1(sun)
             modes["sun"] = sun
-        
+
         if self.include_time:
             time = [x[k] for k in ["date_sin", "date_cos", "time_sin", "time_cos"]]
             time = torch.cat(time, dim=1).float()

pvnet/models/late_fusion/site_encoders/encoders.py

@@ -1,6 +1,4 @@
-"""Encoder modules for the site-level PV data.
-
-"""
+"""Encoder modules for the site-level PV data."""
 
 import einops
 import torch
@@ -11,6 +9,7 @@ from pvnet.models.late_fusion.linear_networks.networks import ResFCNet
 from pvnet.models.late_fusion.site_encoders.basic_blocks import AbstractSitesEncoder
 
 
+# TODO update this to work with the new sample data format
 class SimpleLearnedAggregator(AbstractSitesEncoder):
     """A simple model which learns a different weighted-average across all PV sites for each GSP.
 
@@ -127,8 +126,7 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
         kv_res_block_layers: int = 2,
         use_id_in_value: bool = False,
         target_id_dim: int = 318,
-        target_key_to_use: str = "gsp",
-        input_key_to_use: str = "site",
+        key_to_use: str = "generation",
         num_channels: int = 1,
         num_sites_in_inference: int = 1,
     ):
@@ -149,8 +147,7 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
             use_id_in_value: Whether to use a site ID embedding in network used to produce the
                 value for the attention layer.
             target_id_dim: The number of unique IDs.
-            target_key_to_use: The key to use for the target in the attention layer.
-            input_key_to_use: The key to use for the input in the attention layer.
+            key_to_use: The key to use in the attention layer.
             num_channels: Number of channels in the input data
             num_sites_in_inference: Number of sites to use in inference.
                 This is used to determine the number of sites to use in the
@@ -161,11 +158,10 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
         super().__init__(sequence_length, num_sites, out_features)
         self.sequence_length = sequence_length
         self.target_id_embedding = nn.Embedding(target_id_dim, out_features)
-        self.site_id_embedding = nn.Embedding(num_sites, id_embed_dim)
+        self.id_embedding = nn.Embedding(num_sites, id_embed_dim)
         self._ids = nn.parameter.Parameter(torch.arange(num_sites), requires_grad=False)
         self.use_id_in_value = use_id_in_value
-        self.target_key_to_use = target_key_to_use
-        self.input_key_to_use = input_key_to_use
+        self.key_to_use = key_to_use
         self.num_channels = num_channels
         self.num_sites_in_inference = num_sites_in_inference
 
@@ -206,7 +202,7 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
     def _encode_inputs(self, x: TensorBatch) -> tuple[torch.Tensor, int]:
         # Shape: [batch size, sequence length, number of sites]
         # Shape: [batch size,  station_id, sequence length,  channels]
-        input_data = x[f"{self.input_key_to_use}"]
+        input_data = x[f"{self.key_to_use}"]
         if len(input_data.shape) == 2:  # one site per sample
             input_data = input_data.unsqueeze(-1)  # add dimension of 1 to end to make 3D
         if len(input_data.shape) == 4:  # Has multiple channels
@@ -216,16 +212,11 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
             input_data = input_data[:, : self.sequence_length]
         site_seqs = input_data.float()
         batch_size = site_seqs.shape[0]
-        site_seqs = site_seqs.swapaxes(1, 2)  # [batch size, Site ID, sequence length]
+        site_seqs = site_seqs.swapaxes(1, 2)  # [batch size, location ID, sequence length]
         return site_seqs, batch_size
 
     def _encode_query(self, x: TensorBatch) -> torch.Tensor:
-        if self.target_key_to_use == "gsp":
-            # GSP seems to have a different structure
-            ids = x[f"{self.target_key_to_use}_id"]
-        else:
-            ids = x[f"{self.input_key_to_use}_id"]
-        ids = ids.int()
+        ids = x["location_id"].int()
         query = self.target_id_embedding(ids).unsqueeze(1)
         return query
 
@@ -233,9 +224,9 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
         site_seqs, batch_size = self._encode_inputs(x)
 
         # site ID embeddings are the same for each sample
-        site_id_embed = torch.tile(self.site_id_embedding(self._ids), (batch_size, 1, 1))
+        id_embed = torch.tile(self.id_embedding(self._ids), (batch_size, 1, 1))
         # Each concated (site sequence, site ID embedding) is processed with encoder
-        x_seq_in = torch.cat((site_seqs, site_id_embed), dim=2).flatten(0, 1)
+        x_seq_in = torch.cat((site_seqs, id_embed), dim=2).flatten(0, 1)
         key = self._key_encoder(x_seq_in)
 
         # Reshape to [batch size, site, kdim]
@@ -247,9 +238,9 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
 
         if self.use_id_in_value:
             # site ID embeddings are the same for each sample
-            site_id_embed = torch.tile(self.value_id_embedding(self._ids), (batch_size, 1, 1))
+            id_embed = torch.tile(self.value_id_embedding(self._ids), (batch_size, 1, 1))
             # Each concated (site sequence, site ID embedding) is processed with encoder
-            x_seq_in = torch.cat((site_seqs, site_id_embed), dim=2).flatten(0, 1)
+            x_seq_in = torch.cat((site_seqs, id_embed), dim=2).flatten(0, 1)
         else:
             # Encode each site sequence independently
             x_seq_in = site_seqs.flatten(0, 1)
@@ -260,9 +251,8 @@ class SingleAttentionNetwork(AbstractSitesEncoder):
         return value
 
     def _attention_forward(
-        self, x: dict, 
-        average_attn_weights: bool = True
-    ) -> tuple[torch.Tensor, torch.Tensor:]:
+        self, x: dict, average_attn_weights: bool = True
+    ) -> tuple[torch.Tensor, torch.Tensor :]:
         query = self._encode_query(x)
         key = self._encode_key(x)
         value = self._encode_value(x)

pvnet/training/lightning_module.py

@@ -45,9 +45,9 @@ class PVNetLightningModule(pl.LightningModule):
         self.lr = None
 
     def transfer_batch_to_device(
-        self, 
-        batch: TensorBatch, 
-        device: torch.device, 
+        self,
+        batch: TensorBatch,
+        device: torch.device,
         dataloader_idx: int,
     ) -> dict:
         """Method to move custom batches to a given device"""
@@ -75,7 +75,7 @@ class PVNetLightningModule(pl.LightningModule):
         losses = 2 * torch.cat(losses, dim=2)
 
         return losses.mean()
-    
+
     def configure_optimizers(self):
         """Configure the optimizers using learning rate found with LR finder if used"""
         if self.lr is not None:
@@ -84,7 +84,7 @@ class PVNetLightningModule(pl.LightningModule):
         return self._optimizer(self.model)
 
     def _calculate_common_losses(
-        self, 
+        self,
         y: torch.Tensor,
         y_hat: torch.Tensor,
     ) -> dict[str, torch.Tensor]:
@@ -96,15 +96,15 @@ class PVNetLightningModule(pl.LightningModule):
             losses["quantile_loss"] = self._calculate_quantile_loss(y_hat, y)
             y_hat = self.model._quantiles_to_prediction(y_hat)
 
-        losses.update({"MSE":  F.mse_loss(y_hat, y), "MAE": F.l1_loss(y_hat, y)})
+        losses.update({"MSE": F.mse_loss(y_hat, y), "MAE": F.l1_loss(y_hat, y)})
 
         return losses
-    
+
     def training_step(self, batch: TensorBatch, batch_idx: int) -> torch.Tensor:
         """Run training step"""
         y_hat = self.model(batch)
 
-        y = batch[self.model._target_key][:, -self.model.forecast_len :]
+        y = batch["generation"][:, -self.model.forecast_len :]
 
         losses = self._calculate_common_losses(y, y_hat)
         losses = {f"{k}/train": v for k, v in losses.items()}
@@ -116,10 +116,10 @@ class PVNetLightningModule(pl.LightningModule):
         else:
             opt_target = losses["MAE/train"]
         return opt_target
-    
+
     def _calculate_val_losses(
-        self, 
-        y: torch.Tensor, 
+        self,
+        y: torch.Tensor,
         y_hat: torch.Tensor,
     ) -> dict[str, torch.Tensor]:
         """Calculate additional losses only run in validation"""
@@ -138,28 +138,25 @@ class PVNetLightningModule(pl.LightningModule):
         return losses
 
     def _calculate_step_metrics(
-        self, 
-        y: torch.Tensor, 
-        y_hat: torch.Tensor, 
+        self,
+        y: torch.Tensor,
+        y_hat: torch.Tensor,
     ) -> tuple[np.array, np.array]:
         """Calculate the MAE and MSE at each forecast step"""
 
         mae_each_step = torch.mean(torch.abs(y_hat - y), dim=0).cpu().numpy()
         mse_each_step = torch.mean((y_hat - y) ** 2, dim=0).cpu().numpy()
-       
+
         return mae_each_step, mse_each_step
-    
+
     def _store_val_predictions(self, batch: TensorBatch, y_hat: torch.Tensor) -> None:
         """Internally store the validation predictions"""
-        
-        target_key = self.model._target_key
 
-        y = batch[target_key][:, -self.model.forecast_len :].cpu().numpy()
-        y_hat = y_hat.cpu().numpy() 
-        ids = batch[f"{target_key}_id"].cpu().numpy()
+        y = batch["generation"][:, -self.model.forecast_len :].cpu().numpy()
+        y_hat = y_hat.cpu().numpy()
+        ids = batch["location_id"].cpu().numpy()
         init_times_utc = pd.to_datetime(
-            batch[f"{target_key}_time_utc"][:, self.model.history_len+1]
-            .cpu().numpy().astype("datetime64[ns]")
+            batch["time_utc"][:, self.model.history_len + 1].cpu().numpy().astype("datetime64[ns]")
         )
 
         if self.model.use_quantile_regression:
@@ -170,7 +167,7 @@ class PVNetLightningModule(pl.LightningModule):
 
         ds_preds_batch = xr.Dataset(
             data_vars=dict(
-                y_hat=(["sample_num", "forecast_step",  "p_level"], y_hat),
+                y_hat=(["sample_num", "forecast_step", "p_level"], y_hat),
                 y=(["sample_num", "forecast_step"], y),
             ),
             coords=dict(
@@ -186,7 +183,7 @@ class PVNetLightningModule(pl.LightningModule):
         # Set up stores which we will fill during validation
         self.all_val_results: list[xr.Dataset] = []
         self._val_horizon_maes: list[np.array] = []
-        if self.current_epoch==0:
+        if self.current_epoch == 0:
             self._val_persistence_horizon_maes: list[np.array] = []
 
         # Plot some sample forecasts
@@ -197,9 +194,9 @@ class PVNetLightningModule(pl.LightningModule):
 
         for plot_num in range(num_figures):
             idxs = np.arange(plots_per_figure) + plot_num * plots_per_figure
-            idxs = idxs[idxs<len(val_dataset)]
+            idxs = idxs[idxs < len(val_dataset)]
 
-            if len(idxs)==0:
+            if len(idxs) == 0:
                 continue
 
             batch = collate_fn([val_dataset[i] for i in idxs])
@@ -207,19 +204,16 @@ class PVNetLightningModule(pl.LightningModule):
 
             # Batch validation check only during sanity check phase - use first batch
             if self.trainer.sanity_checking and plot_num == 0:
-                validate_batch_against_config(
-                    batch=batch,
-                    model=self.model
-                )
-            
+                validate_batch_against_config(batch=batch, model=self.model)
+
             with torch.no_grad():
                 y_hat = self.model(batch)
-            
+
             fig = plot_sample_forecasts(
                 batch,
                 y_hat,
                 quantiles=self.model.output_quantiles,
-                key_to_plot=self.model._target_key,
+                key_to_plot="generation",
             )
 
             plot_name = f"val_forecast_samples/sample_set_{plot_num}"
@@ -238,7 +232,7 @@ class PVNetLightningModule(pl.LightningModule):
         # Internally store the val predictions
         self._store_val_predictions(batch, y_hat)
 
-        y = batch[self.model._target_key][:, -self.model.forecast_len :]
+        y = batch["generation"][:, -self.model.forecast_len :]
 
         losses = self._calculate_common_losses(y, y_hat)
         losses = {f"{k}/val": v for k, v in losses.items()}
@@ -262,21 +256,22 @@ class PVNetLightningModule(pl.LightningModule):
 
         # Calculate the persistance losses - we only need to do this once per training run
         # not every epoch
-        if self.current_epoch==0:
+        if self.current_epoch == 0:
             y_persist = (
-                batch[self.model._target_key][:, -(self.model.forecast_len+1)]
-                .unsqueeze(1).expand(-1, self.model.forecast_len)
+                batch["generation"][:, -(self.model.forecast_len + 1)]
+                .unsqueeze(1)
+                .expand(-1, self.model.forecast_len)
             )
             mae_step_persist, mse_step_persist = self._calculate_step_metrics(y, y_persist)
             self._val_persistence_horizon_maes.append(mae_step_persist)
             losses.update(
                 {
-                    "MAE/val_persistence": mae_step_persist.mean(), 
-                    "MSE/val_persistence": mse_step_persist.mean()
+                    "MAE/val_persistence": mae_step_persist.mean(),
+                    "MSE/val_persistence": mse_step_persist.mean(),
                 }
             )
 
-        # Log the metrics
+        # Log the metrics
         self.log_dict(losses, on_step=False, on_epoch=True)
 
     def on_validation_epoch_end(self) -> None:
@@ -289,7 +284,7 @@ class PVNetLightningModule(pl.LightningModule):
         self._val_horizon_maes = []
 
         # We only run this on the first epoch
-        if self.current_epoch==0:
+        if self.current_epoch == 0:
             val_persistence_horizon_maes = np.mean(self._val_persistence_horizon_maes, axis=0)
             self._val_persistence_horizon_maes = []
 
@@ -321,25 +316,25 @@ class PVNetLightningModule(pl.LightningModule):
                 wandb_log_dir = self.logger.experiment.dir
                 filepath = f"{wandb_log_dir}/validation_results.netcdf"
                 ds_val_results.to_netcdf(filepath)
-                
-                # Uplodad to wandb
+
+                # Uplodad to wandb
                 self.logger.experiment.save(filepath, base_path=wandb_log_dir, policy="now")
-            
+
             # Create the horizon accuracy curve
             horizon_mae_plot = wandb_line_plot(
-                x=np.arange(self.model.forecast_len), 
+                x=np.arange(self.model.forecast_len),
                 y=val_horizon_maes,
                 xlabel="Horizon step",
                 ylabel="MAE",
                 title="Val horizon loss curve",
             )
-            
+
             wandb.log({"val_horizon_mae_plot": horizon_mae_plot})
 
             # Create persistence horizon accuracy curve but only on first epoch
-            if self.current_epoch==0:
+            if self.current_epoch == 0:
                 persist_horizon_mae_plot = wandb_line_plot(
-                    x=np.arange(self.model.forecast_len), 
+                    x=np.arange(self.model.forecast_len),
                     y=val_persistence_horizon_maes,
                     xlabel="Horizon step",
                     ylabel="MAE",

pvnet/training/plots.py

@@ -32,9 +32,9 @@ def plot_sample_forecasts(
 
     y = batch[key_to_plot].cpu().numpy()
     y_hat = y_hat.cpu().numpy()    
-    ids = batch[f"{key_to_plot}_id"].cpu().numpy().squeeze()
+    ids = batch["location_id"].cpu().numpy().squeeze()
     times_utc = pd.to_datetime(
-        batch[f"{key_to_plot}_time_utc"].cpu().numpy().squeeze().astype("datetime64[ns]")
+        batch["time_utc"].cpu().numpy().squeeze().astype("datetime64[ns]")
     )
     batch_size = y.shape[0]
 

pvnet/utils.py

@@ -1,4 +1,5 @@
 """Utils"""
+
 import logging
 from typing import TYPE_CHECKING
 
@@ -17,7 +18,7 @@ PYTORCH_WEIGHTS_NAME = "model_weights.safetensors"
 MODEL_CONFIG_NAME = "model_config.yaml"
 DATA_CONFIG_NAME = "data_config.yaml"
 DATAMODULE_CONFIG_NAME = "datamodule_config.yaml"
-FULL_CONFIG_NAME =  "full_experiment_config.yaml"
+FULL_CONFIG_NAME = "full_experiment_config.yaml"
 MODEL_CARD_NAME = "README.md"
 
 
@@ -93,37 +94,41 @@ def print_config(
 
 
 def validate_batch_against_config(
-    batch: dict, 
+    batch: dict,
     model: "BaseModel",
 ) -> None:
     """Validates tensor shapes in batch against model configuration."""
     logger.info("Performing batch shape validation against model config.")
-    
+
     # NWP validation
-    if hasattr(model, 'nwp_encoders_dict'):
+    if hasattr(model, "nwp_encoders_dict"):
         if "nwp" not in batch:
             raise ValueError(
                 "Model configured with 'nwp_encoders_dict' but 'nwp' data missing from batch."
             )
-            
+
         for source, nwp_data in batch["nwp"].items():
             if source in model.nwp_encoders_dict:
-
-                enc = model.nwp_encoders_dict[source]                
+                enc = model.nwp_encoders_dict[source]
                 expected_channels = enc.in_channels
                 if model.add_image_embedding_channel:
                     expected_channels -= 1
 
-                expected = (nwp_data["nwp"].shape[0], enc.sequence_length, 
-                           expected_channels, enc.image_size_pixels, enc.image_size_pixels)
+                expected = (
+                    nwp_data["nwp"].shape[0],
+                    enc.sequence_length,
+                    expected_channels,
+                    enc.image_size_pixels,
+                    enc.image_size_pixels,
+                )
                 if tuple(nwp_data["nwp"].shape) != expected:
-                    actual_shape = tuple(nwp_data['nwp'].shape)
+                    actual_shape = tuple(nwp_data["nwp"].shape)
                     raise ValueError(
                         f"NWP.{source} shape mismatch: expected {expected}, got {actual_shape}"
                     )
 
     # Satellite validation
-    if hasattr(model, 'sat_encoder'):
+    if hasattr(model, "sat_encoder"):
         if "satellite_actual" not in batch:
             raise ValueError(
                 "Model configured with 'sat_encoder' but 'satellite_actual' missing from batch."
@@ -134,14 +139,19 @@ def validate_batch_against_config(
         if model.add_image_embedding_channel:
             expected_channels -= 1
 
-        expected = (batch["satellite_actual"].shape[0], enc.sequence_length, expected_channels, 
-                enc.image_size_pixels, enc.image_size_pixels)
+        expected = (
+            batch["satellite_actual"].shape[0],
+            enc.sequence_length,
+            expected_channels,
+            enc.image_size_pixels,
+            enc.image_size_pixels,
+        )
         if tuple(batch["satellite_actual"].shape) != expected:
-            actual_shape = tuple(batch['satellite_actual'].shape)
+            actual_shape = tuple(batch["satellite_actual"].shape)
             raise ValueError(f"Satellite shape mismatch: expected {expected}, got {actual_shape}")
 
-    # GSP/Site validation
-    key = model._target_key
+    # generation validation
+    key = "generation"
     if key in batch:
         total_minutes = model.history_minutes + model.forecast_minutes
         interval = model.interval_minutes

{pvnet-5.2.3.dist-info → pvnet-5.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PVNet
-Version: 5.2.3
+Version: 5.3.0
 Summary: PVNet
 Author-email: Peter Dudfield <info@openclimatefix.org>
 Requires-Python: <3.14,>=3.11

{pvnet-5.2.3.dist-info → pvnet-5.3.0.dist-info}/RECORD

@@ -1,14 +1,14 @@
 pvnet/__init__.py,sha256=TAZm88TJ5ieL1XjEyRg1LciIGuSScEucdAruQLfM92I,25
-pvnet/datamodule.py,sha256=sTACPJXPaqojpxf86wldqxxlnFRoPvlvRHkmcGSsmDw,6368
+pvnet/datamodule.py,sha256=wc1RQfFhgW9Hxyw7vrpFERhOd2FmjDsO1x49J2erOYk,5750
 pvnet/load_model.py,sha256=P1QODX_mJRnKZ_kIll9BlOjK_A1W4YM3QG-mZd-2Mcc,3852
 pvnet/optimizers.py,sha256=1N4b-Xd6QiIrcUU8cbU326bbFC0BvMNIV8VYWtGILJc,6548
-pvnet/utils.py,sha256=6hVKQN8F89pJbiC9VSuHCm5yJqzIzs7hLF3ztkBU-TY,5895
+pvnet/utils.py,sha256=L3MDF5m1Ez_btAZZ8t-T5wXLzFmyj7UZtorA91DEpFw,6003
 pvnet/models/__init__.py,sha256=owzZ9xkD0DRTT51mT2Dx_p96oJjwDz57xo_MaMIEosk,145
-pvnet/models/base_model.py,sha256=CnQaaf2kAdOcXqo1319nWa120mHfLQiwOQ639m4OzPk,16182
-pvnet/models/ensemble.py,sha256=1mFUEsl33kWcLL5d7zfDm9ypWxgAxBHgBiJLt0vwTeg,2363
+pvnet/models/base_model.py,sha256=V-vBqtzZc_c8Ho5hVo_ikq2wzZ7hsAIM7I4vhzGDfNc,16051
+pvnet/models/ensemble.py,sha256=USpNQ0O5eiffapLPE9T6gR-uK9f_3E4pX3DK7Lmkn2U,2228
 pvnet/models/late_fusion/__init__.py,sha256=Jf0B-E0_5IvSBFoj1wvnPtwYDxs4pRIFm5qHv--Bbps,26
 pvnet/models/late_fusion/basic_blocks.py,sha256=_cYGVyAIyEJS4wd-DEAXQXu0br66guZJn3ugoebWqZ0,1479
-pvnet/models/late_fusion/late_fusion.py,sha256=7uQPo_OlNXrJOE9nYHTEvwJx2POKg4drJfdnPxwiaJU,16283
+pvnet/models/late_fusion/late_fusion.py,sha256=kQUnyqMykmwc0GdoFhNXYStJPrjr3hFSvUNe8FumVx4,15260
 pvnet/models/late_fusion/encoders/__init__.py,sha256=bLBQdnCeLYhwISW0t88ZZBz-ebS94m7ZwBcsofWMHR4,51
 pvnet/models/late_fusion/encoders/basic_blocks.py,sha256=DGkFFIZv4S4FLTaAIOrAngAFBpgZQHfkGM4dzezZLk4,3044
 pvnet/models/late_fusion/encoders/encoders3d.py,sha256=9fmqVHO73F-jN62w065cgEQI_icNFC2nQH6ZEGvTHxU,7116
@@ -17,13 +17,13 @@ pvnet/models/late_fusion/linear_networks/basic_blocks.py,sha256=RnwdeuX_-itY4ncM
 pvnet/models/late_fusion/linear_networks/networks.py,sha256=exEIz_Z85f8nSwcvp4wqiiLECEAg9YbkKhSZJvFy75M,2231
 pvnet/models/late_fusion/site_encoders/__init__.py,sha256=QoUiiWWFf12vEpdkw0gO4TWpOEoI_tgAyUFCWFFpYAk,45
 pvnet/models/late_fusion/site_encoders/basic_blocks.py,sha256=iEB_N7ZL5HMQ1hZM6H32A71GCwP7YbErUx0oQF21PQM,1042
-pvnet/models/late_fusion/site_encoders/encoders.py,sha256=k4z690cfcP6J4pm2KtDujHN-W3uOl7QY0WvBIu1tM8c,11703
+pvnet/models/late_fusion/site_encoders/encoders.py,sha256=DcTV2LeZ0pSZpGFmsmPEqYIhmPQeYCNi3UM406zHm14,11310
 pvnet/training/__init__.py,sha256=FKxmPZ59Vuj5_mXomN4saJ3En5M-aDMxSs6OttTQOcg,49
-pvnet/training/lightning_module.py,sha256=KcEbHYBe_Gx0as0-A7bggoMjev-A_i6Y3PHGRaYllTg,12956
-pvnet/training/plots.py,sha256=4xID7TBA4IazaARaCN5AoG5fFPJF1wIprn0y6I0C31c,2469
+pvnet/training/lightning_module.py,sha256=57sT7bPCU7mJw4EskzOE-JJ9JhWIuAbs40_x5RoBbA8,12705
+pvnet/training/plots.py,sha256=7JtjA9zIotuoKZ2l0fbS-FZDB48TcIk_-XLA2EWVMv4,2448
 pvnet/training/train.py,sha256=Sry2wYgggUmtIB-k_umFts7xMr2roEL76NCu9ySbLUY,4107
-pvnet-5.2.3.dist-info/licenses/LICENSE,sha256=tKUnlSmcLBWMJWkHx3UjZGdrjs9LidGwLo0jsBUBAwU,1077
-pvnet-5.2.3.dist-info/METADATA,sha256=xA54YM0qDAlvtNSEQXwLLYF9S6sgo-kejaB0awBb0MA,16479
-pvnet-5.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pvnet-5.2.3.dist-info/top_level.txt,sha256=4mg6WjeW05SR7pg3-Q4JRE2yAoutHYpspOsiUzYVNv0,6
-pvnet-5.2.3.dist-info/RECORD,,
+pvnet-5.3.0.dist-info/licenses/LICENSE,sha256=tKUnlSmcLBWMJWkHx3UjZGdrjs9LidGwLo0jsBUBAwU,1077
+pvnet-5.3.0.dist-info/METADATA,sha256=b4Ki0jGoNNEd1VopMvR5p-iasCi0ZVtGwA-RfoHRCWw,16479
+pvnet-5.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pvnet-5.3.0.dist-info/top_level.txt,sha256=4mg6WjeW05SR7pg3-Q4JRE2yAoutHYpspOsiUzYVNv0,6
+pvnet-5.3.0.dist-info/RECORD,,