Add quantile-quantile plot to wandb custom charts

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(

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,

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,