PVNet_summation 1.1.2__py3-none-any.whl

pvnet_summation/optimizers.py

@@ -0,0 +1,219 @@
+"""Optimizer factory-function classes.
+"""
+
+from abc import ABC, abstractmethod
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+
+def find_submodule_parameters(model: Module, search_modules: list[Module]) -> list[Parameter]:
+    """Finds all parameters within given submodule types
+
+    Args:
+        model: torch Module to search through
+        search_modules: List of submodule types to search for
+    """
+    if isinstance(model, search_modules):
+        return model.parameters()
+
+    children = list(model.children())
+    if len(children) == 0:
+        return []
+    else:
+        params = []
+        for c in children:
+            params += find_submodule_parameters(c, search_modules)
+        return params
+
+
+def find_other_than_submodule_parameters(
+    model: Module, 
+    ignore_modules: list[Module],
+) -> list[Parameter]:
+    """Finds all parameters not with given submodule types
+
+    Args:
+        model: torch Module to search through
+        ignore_modules: List of submodule types to ignore
+    """
+    if isinstance(model, ignore_modules):
+        return []
+
+    children = list(model.children())
+    if len(children) == 0:
+        return model.parameters()
+    else:
+        params = []
+        for c in children:
+            params += find_other_than_submodule_parameters(c, ignore_modules)
+        return params
+
+
+class AbstractOptimizer(ABC):
+    """Abstract class for optimizer
+
+    Optimizer classes will be used by model like:
+    > OptimizerGenerator = AbstractOptimizer()
+    > optimizer = OptimizerGenerator(model)
+    The returned object `optimizer` must be something that may be returned by `pytorch_lightning`'s
+    `configure_optimizers()` method.
+    See :
+        https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#configure-optimizers
+
+    """
+
+    @abstractmethod
+    def __call__(self):
+        """Abstract call"""
+        pass
+
+
+class Adam(AbstractOptimizer):
+    """Adam optimizer"""
+
+    def __init__(self, lr: float = 0.0005, **kwargs):
+        """Adam optimizer"""
+        self.lr = lr
+        self.kwargs = kwargs
+
+    def __call__(self, model: Module):
+        """Return optimizer"""
+        return torch.optim.Adam(model.parameters(), lr=self.lr, **self.kwargs)
+
+
+class AdamW(AbstractOptimizer):
+    """AdamW optimizer"""
+
+    def __init__(self, lr: float = 0.0005, **kwargs):
+        """AdamW optimizer"""
+        self.lr = lr
+        self.kwargs = kwargs
+
+    def __call__(self, model: Module):
+        """Return optimizer"""
+        return torch.optim.AdamW(model.parameters(), lr=self.lr, **self.kwargs)
+
+
+
+class EmbAdamWReduceLROnPlateau(AbstractOptimizer):
+    """AdamW optimizer and reduce on plateau scheduler"""
+
+    def __init__(
+        self, 
+        lr: float = 0.0005, 
+        weight_decay: float = 0.01, 
+        patience: int = 3, 
+        factor: float = 0.5, 
+        threshold: float = 2e-4, 
+        **opt_kwargs,
+    ):
+        """AdamW optimizer and reduce on plateau scheduler"""
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.patience = patience
+        self.factor = factor
+        self.threshold = threshold
+        self.opt_kwargs = opt_kwargs
+
+    def __call__(self, model):
+        """Return optimizer"""
+
+        search_modules = (torch.nn.Embedding,)
+
+        no_decay = find_submodule_parameters(model, search_modules)
+        decay = find_other_than_submodule_parameters(model, search_modules)
+
+        optim_groups = [
+            {"params": decay, "weight_decay": self.weight_decay},
+            {"params": no_decay, "weight_decay": 0.0},
+        ]
+        opt = torch.optim.AdamW(optim_groups, lr=self.lr, **self.opt_kwargs)
+
+        sch = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            opt,
+            factor=self.factor,
+            patience=self.patience,
+            threshold=self.threshold,
+        )
+        sch = {
+            "scheduler": sch,
+            "monitor": "quantile_loss/val" if model.use_quantile_regression else "MAE/val",
+        }
+        return [opt], [sch]
+
+
+class AdamWReduceLROnPlateau(AbstractOptimizer):
+    """AdamW optimizer and reduce on plateau scheduler"""
+
+    def __init__(
+        self, 
+        lr: float = 0.0005, 
+        patience: int = 3, 
+        factor: float = 0.5, 
+        threshold: float = 2e-4, 
+        step_freq=None, 
+        **opt_kwargs,
+    ):
+        """AdamW optimizer and reduce on plateau scheduler"""
+        self.lr = lr
+        self.patience = patience
+        self.factor = factor
+        self.threshold = threshold
+        self.step_freq = step_freq
+        self.opt_kwargs = opt_kwargs
+
+    def _call_multi(self, model):
+        remaining_params = {k: p for k, p in model.named_parameters()}
+
+        group_args = []
+
+        for key in self.lr.keys():
+            if key == "default":
+                continue
+
+            submodule_params = []
+            for param_name in list(remaining_params.keys()):
+                if param_name.startswith(key):
+                    submodule_params += [remaining_params.pop(param_name)]
+
+            group_args += [{"params": submodule_params, "lr": self.lr[key]}]
+
+        remaining_params = [p for k, p in remaining_params.items()]
+        group_args += [{"params": remaining_params}]
+
+        opt = torch.optim.AdamW(
+            group_args, 
+            lr=self.lr["default"] if model.lr is None else model.lr, 
+            **self.opt_kwargs,
+        )
+        sch = {
+            "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
+                opt,
+                factor=self.factor,
+                patience=self.patience,
+                threshold=self.threshold,
+            ),
+            "monitor": "quantile_loss/val" if model.use_quantile_regression else "MAE/val",
+        }
+
+        return [opt], [sch]
+
+    def __call__(self, model):
+        """Return optimizer"""
+        if not isinstance(self.lr, float):
+            return self._call_multi(model)
+        else:
+            opt = torch.optim.AdamW(model.parameters(), lr=self.lr, **self.opt_kwargs)
+            sch = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                opt,
+                factor=self.factor,
+                patience=self.patience,
+                threshold=self.threshold,
+            )
+            sch = {
+                "scheduler": sch,
+                "monitor": "quantile_loss/val" if model.use_quantile_regression else "MAE/val",
+            }
+            return [opt], [sch]

pvnet_summation/training/__init__.py

@@ -0,0 +1,3 @@
+"""Training submodule"""
+from .lightning_module import PVNetSummationLightningModule
+from .train import train

pvnet_summation/training/lightning_module.py

@@ -0,0 +1,278 @@
+"""Pytorch lightning module for training PVNet models"""
+
+import lightning.pytorch as pl
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+import wandb
+from ocf_data_sampler.numpy_sample.common_types import TensorBatch
+from torch.utils.data import default_collate
+
+from pvnet_summation.models.base_model import BaseModel
+from pvnet_summation.optimizers import AbstractOptimizer
+from pvnet_summation.training.plots import plot_sample_forecasts, wandb_line_plot
+
+
+class PVNetSummationLightningModule(pl.LightningModule):
+    """Lightning module for training PVNet models"""
+
+    def __init__(
+        self,
+        model: BaseModel,
+        optimizer: AbstractOptimizer,
+    ):
+        """Lightning module for training PVNet models
+
+        Args:
+            model: The PVNet model
+            optimizer: Optimizer
+        """
+        super().__init__()
+
+        self.model = model
+        self._optimizer = optimizer
+
+        # Model must have lr to allow tuning
+        # This setting is only used when lr is tuned with callback
+        self.lr = None
+
+
+    def _calculate_quantile_loss(self, y_quantiles: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """Calculate quantile loss.
+
+        Note:
+            Implementation copied from:
+                https://pytorch-forecasting.readthedocs.io/en/stable/_modules/pytorch_forecasting
+                /metrics/quantile.html#QuantileLoss.loss
+
+        Args:
+            y_quantiles: Quantile prediction of network
+            y: Target values
+
+        Returns:
+            Quantile loss
+        """
+        losses = []
+        for i, q in enumerate(self.model.output_quantiles):
+            errors = y - y_quantiles[..., i]
+            losses.append(torch.max((q - 1) * errors, q * errors).unsqueeze(-1))
+        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:
+            # Use learning rate found by learning rate finder callback
+            self._optimizer.lr = self.lr
+        return self._optimizer(self.model)
+
+    def _calculate_common_losses(
+        self, 
+        y: torch.Tensor,
+        y_hat: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        """Calculate losses common to train, and val"""
+
+        losses = {}
+
+        if self.model.use_quantile_regression:
+            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)})
+
+        return losses
+    
+    def training_step(self, batch: TensorBatch, batch_idx: int) -> torch.Tensor:
+        """Run training step"""
+
+        y_hat = self.model(batch)
+
+        y = batch["target"]
+
+        losses = self._calculate_common_losses(y, y_hat)
+        losses = {f"{k}/train": v for k, v in losses.items()}
+
+        self.log_dict(losses, on_step=True, on_epoch=True)
+
+        if self.model.use_quantile_regression:
+            opt_target = losses["quantile_loss/train"]
+        else:
+            opt_target = losses["MAE/train"]
+        return opt_target
+    
+    def _calculate_val_losses(
+        self, 
+        y: torch.Tensor, 
+        y_hat: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        """Calculate additional losses only run in validation"""
+
+        losses = {}
+
+        if self.model.use_quantile_regression:
+            metric_name = "val_fraction_below/fraction_below_{:.2f}_quantile"
+            # Add fraction below each quantile for calibration
+            for i, quantile in enumerate(self.model.output_quantiles):
+                below_quant = y <= y_hat[..., i]
+                # Mask values small values, which are dominated by night
+                mask = y >= 0.01
+                losses[metric_name.format(quantile)] = below_quant[mask].float().mean()
+
+        return losses
+
+    def _calculate_step_metrics(
+        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 on_validation_epoch_start(self):
+        """Run at start of val period"""
+        # Set up stores which we will fill during validation
+        self._val_horizon_maes: list[np.array] = []
+        if self.current_epoch==0:
+            self._val_persistence_horizon_maes: list[np.array] = []
+            self._val_loc_sum_horizon_maes: list[np.array] = []
+        
+        # Plot some sample forecasts
+        val_dataset = self.trainer.val_dataloaders.dataset
+
+        plots_per_figure = 16
+        num_figures = 2
+
+        for plot_num in range(num_figures):
+            idxs = np.arange(plots_per_figure) + plot_num * plots_per_figure
+            idxs = idxs[idxs<len(val_dataset)]
+
+            if len(idxs)==0:
+                continue
+
+            batch = default_collate([val_dataset[i] for i in idxs])
+            batch = self.transfer_batch_to_device(batch, self.device, dataloader_idx=0)
+            with torch.no_grad():
+                y_hat = self.model(batch)
+
+            y_loc_sum = self.model.sum_of_locations(batch)
+            
+            fig = plot_sample_forecasts(batch, y_hat, y_loc_sum, self.model.output_quantiles)
+
+            plot_name = f"val_forecast_samples/sample_set_{plot_num}"
+
+            self.logger.experiment.log({plot_name: wandb.Image(fig)})
+
+            plt.close(fig)
+
+    def validation_step(self, batch: TensorBatch, batch_idx: int) -> None:
+        """Run validation step"""
+
+        y_hat = self.model(batch)
+
+        y = batch["target"]
+
+        losses = self._calculate_common_losses(y, y_hat)
+        losses = {f"{k}/val": v for k, v in losses.items()}
+
+        losses.update(self._calculate_val_losses(y, y_hat))
+
+        # Calculate the horizon MAE/MSE metrics
+        if self.model.use_quantile_regression:
+            y_hat_mid = self.model._quantiles_to_prediction(y_hat)
+        else:
+            y_hat_mid = y_hat
+
+        mae_step, mse_step = self._calculate_step_metrics(y, y_hat_mid)
+
+        # Store to make horizon-MAE plot
+        self._val_horizon_maes.append(mae_step)
+
+        # Also add each step to logged metrics
+        losses.update({f"val_step_MAE/step_{i:03}": m for i, m in enumerate(mae_step)})
+        losses.update({f"val_step_MSE/step_{i:03}": m for i, m in enumerate(mse_step)})
+
+        # Calculate the persistence and sum-of-locations losses - we only need to do this once per 
+        # training run not every epoch
+        if self.current_epoch==0:
+            
+            # Persistence
+            y_persist = batch["last_outturn"].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()
+                }
+            )
+
+            # Sum of Locations
+            y_loc_sum = self.model.sum_of_locations(batch)
+            mae_step_loc_sum, mse_step_loc_sum = self._calculate_step_metrics(y, y_loc_sum)
+            self._val_loc_sum_horizon_maes.append(mae_step_loc_sum)
+            losses.update(
+                {
+                    "MAE/val_location_sum": mae_step_loc_sum.mean(), 
+                    "MSE/val_location_sum": mse_step_loc_sum.mean()
+                }
+            )
+
+        # Log the metrics
+        self.log_dict(losses, on_step=False, on_epoch=True)
+
+    def on_validation_epoch_end(self) -> None:
+        """Run on epoch end"""
+
+        val_horizon_maes = np.mean(self._val_horizon_maes, axis=0)
+        self._val_horizon_maes = []
+
+        if isinstance(self.logger, pl.loggers.WandbLogger):
+            
+            # Create the horizon accuracy curve
+            horizon_mae_plot = wandb_line_plot(
+                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 and location-sum horizon accuracy curve on first epoch
+            if self.current_epoch==0:
+                val_persistence_horizon_maes = np.mean(self._val_persistence_horizon_maes, axis=0)
+                del self._val_persistence_horizon_maes
+
+                val_loc_sum_horizon_maes = np.mean(self._val_loc_sum_horizon_maes, axis=0)
+                del self._val_loc_sum_horizon_maes
+
+                persist_horizon_mae_plot = wandb_line_plot(
+                    x=np.arange(self.model.forecast_len), 
+                    y=val_persistence_horizon_maes,
+                    xlabel="Horizon step",
+                    ylabel="MAE",
+                    title="Val persistence horizon loss curve",
+                )
+
+                loc_sum_horizon_mae_plot = wandb_line_plot(
+                    x=np.arange(self.model.forecast_len), 
+                    y=val_loc_sum_horizon_maes,
+                    xlabel="Horizon step",
+                    ylabel="MAE",
+                    title="Val location-sum horizon loss curve",
+                )
+
+                wandb.log(
+                    {
+                        "persistence_val_horizon_mae_plot": persist_horizon_mae_plot,
+                        "location_sum_val_horizon_mae_plot": loc_sum_horizon_mae_plot,
+                    }
+                )

pvnet_summation/training/plots.py

@@ -0,0 +1,91 @@
+"""Plots logged during training"""
+from collections.abc import Sequence
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import pylab
+import torch
+import wandb
+
+from pvnet_summation.data.datamodule import SumTensorBatch
+
+
+def wandb_line_plot(
+    x: Sequence[float], 
+    y: Sequence[float], 
+    xlabel: str, 
+    ylabel: str, 
+    title: str | None = None,
+) -> wandb.plot.CustomChart:
+    """Make a wandb line plot"""
+    data = [[xi, yi] for (xi, yi) in zip(x, y)]
+    table = wandb.Table(data=data, columns=[xlabel, ylabel])
+    return wandb.plot.line(table, xlabel, ylabel, title=title)
+
+
+def plot_sample_forecasts(
+    batch: SumTensorBatch,
+    y_hat: torch.Tensor,
+    y_loc_sum: torch.Tensor,
+    quantiles: list[float] | None,
+) -> plt.Figure:
+    """Plot a batch of data and the forecast from that batch"""
+
+    y = batch["target"].cpu().numpy()
+    y_hat = y_hat.cpu().numpy()
+    y_loc_sum = y_loc_sum.cpu().numpy()
+    times_utc = pd.to_datetime(batch["valid_times"].cpu().numpy().astype("datetime64[ns]"))
+    batch_size = y.shape[0]
+
+    fig, axes = plt.subplots(4, 4, figsize=(16, 16))
+
+    for i, ax in enumerate(axes.ravel()[:batch_size]):
+
+        ax.plot(times_utc[i], y[i], marker=".", color="k", label=r"$y$")
+
+        ax.plot(
+            times_utc[i],
+            y_loc_sum[i], 
+            marker=".",
+            linestyle="-.",
+            color="r",
+            label=r"$\hat{y}_{loc-sum}$",
+        )
+
+        if quantiles is None:
+            ax.plot(
+                times_utc[i],
+                y_hat[i], 
+                marker=".", 
+                color="r", 
+                label=r"$\hat{y}$",
+            )
+        else:
+            cm = pylab.get_cmap("twilight")
+            for nq, q in enumerate(quantiles):
+                ax.plot(
+                    times_utc[i],
+                    y_hat[i, :, nq],
+                    color=cm(q),
+                    label=r"$\hat{y}$" + f"({q})",
+                    alpha=0.7,
+                )
+
+        ax.set_title(f"{times_utc[i][0].date()}", fontsize="small")
+
+        xticks = [t for t in times_utc[i] if t.minute == 0][::2]
+        ax.set_xticks(ticks=xticks, labels=[f"{t.hour:02}" for t in xticks], rotation=90)
+        ax.grid()
+
+    axes[0, 0].legend(loc="best")
+
+    if batch_size<16:
+        for ax in axes.ravel()[batch_size:]:
+            ax.axis("off")
+    
+    for ax in axes[-1, :]:
+        ax.set_xlabel("Time (hour of day)")
+
+    plt.tight_layout()
+
+    return fig