diff --git a/pvnet/training/lightning_module.py b/pvnet/training/lightning_module.py index 18678b78..484ba61b 100644 --- a/pvnet/training/lightning_module.py +++ b/pvnet/training/lightning_module.py @@ -14,7 +14,7 @@ from pvnet.datamodule import collate_fn from pvnet.models.base_model import BaseModel from pvnet.optimizers import AbstractOptimizer -from pvnet.training.plots import plot_sample_forecasts, wandb_line_plot +from pvnet.training.plots import plot_sample_forecasts, wandb_line_plot, wandb_line_plot_custom from pvnet.utils import validate_batch_against_config @@ -129,11 +129,15 @@ def _calculate_val_losses( if self.model.use_quantile_regression: metric_name = "val_fraction_below/fraction_below_{:.2f}_quantile" # Add fraction below each quantile for calibration + val_quantiles = np.zeros(len(self.model.output_quantiles)) 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() + below_quant_masked_mean = below_quant[mask].float().mean() + losses[metric_name.format(quantile)] = below_quant_masked_mean + val_quantiles[i] = below_quant_masked_mean + self._val_quantiles.append(val_quantiles) return losses @@ -185,6 +189,8 @@ def on_validation_epoch_start(self): # 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.model.use_quantile_regression: + self._val_quantiles: list[np.array] = [] if self.current_epoch == 0: self._val_persistence_horizon_maes: list[np.array] = [] @@ -355,6 +361,25 @@ def on_validation_epoch_end(self) -> None: step=self.trainer.global_step, ) + # Create a quantile-quantile plot + if self.model.use_quantile_regression: + val_quantiles = np.mean(self._val_quantiles, axis=0) + self._val_quantiles = [] + + qq_plot = wandb_line_plot( + x=self.model.output_quantiles, + y=val_quantiles, + xlabel="True quantiles", + ylabel="Predicted quantiles", + title="Quantile-quantile plot", + add_identity_line=True, + ) + + wandb.log( + {"quantile_quantile": qq_plot}, + step=self.trainer.global_step, + ) + # Create persistence horizon accuracy curve but only on first epoch if self.current_epoch == 0: persist_horizon_mae_plot = wandb_line_plot( diff --git a/pvnet/training/plots.py b/pvnet/training/plots.py index 21121a37..3e37f90a 100644 --- a/pvnet/training/plots.py +++ b/pvnet/training/plots.py @@ -14,13 +14,31 @@ def wandb_line_plot( y: Sequence[float], xlabel: str, ylabel: str, - title: str | None = None + title: str | None = None, + add_identity_line: bool = False, ) -> 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) - + # Main series data + data = [[xi, yi, "Data"] for xi, yi in zip(x, y)] + + # Add identity line endpoints if requested + if add_identity_line: + min_val, max_val = min(x), max(x) + data.append([min_val, min_val, "x=y"]) + data.append([max_val, max_val, "x=y"]) + + table = wandb.Table(data=data, columns=[xlabel, ylabel, "Series"]) + + # stroke=None creates a clean single line; stroke="Series" creates multi-line legend + stroke_col = "Series" if add_identity_line else None + + return wandb.plot.line( + table=table, + x=xlabel, + y=ylabel, + stroke=stroke_col, + title=title + ) def plot_sample_forecasts( batch: TensorBatch, diff --git a/tests/training/test_train.py b/tests/training/test_train.py index 1760604d..c3f0e47e 100644 --- a/tests/training/test_train.py +++ b/tests/training/test_train.py @@ -20,7 +20,7 @@ def trainer_cfg_cpu() -> dict: """Tiny CPU-only Trainer config.""" return { "_target_": "lightning.pytorch.Trainer", - "max_epochs": 1, + "max_epochs": 2, "limit_train_batches": 1, "limit_val_batches": 1, "accelerator": "cpu", @@ -71,6 +71,7 @@ def build_lit_late_fusion_cfg( "_target_": "pvnet.training.lightning_module.PVNetLightningModule", "model": { "_target_": "pvnet.models.LateFusionModel", + "output_quantiles": [0.1, 0.5, 0.9], "sat_encoder": None, "nwp_encoders_dict": None, "add_image_embedding_channel": False,