Make plotting configurable

Author: jacobbieker

pvnet/models/base_model.py

@@ -241,6 +241,7 @@ def __init__(
         output_quantiles: Optional[list[float]] = None,
         target_key: str = "gsp",
         interval_minutes: int = 30,
+        timestep_intervals_to_plot: Optional[list[int]] = None
     ):
         """Abtstract base class for PVNet submodels.
 
@@ -258,6 +259,10 @@ def __init__(
         self._optimizer = optimizer
         self._target_key_name = target_key
         self._target_key = BatchKey[f"{target_key}"]
+        if timestep_intervals_to_plot is not None:
+            for interval in timestep_intervals_to_plot:
+                assert type(interval) in [list, tuple] and len(interval) == 2, ValueError(f"timestep_intervals_to_plot must be a list of tuples or lists of length 2, but got {timestep_intervals_to_plot=}")
+        self.time_step_intervals_to_plot = timestep_intervals_to_plot
 
         # Model must have lr to allow tuning
         # This setting is only used when lr is tuned with callback
@@ -268,12 +273,12 @@ def __init__(
         self.output_quantiles = output_quantiles
 
         # Number of timestemps for 30 minutely data
-        self.history_len_30 = history_minutes // interval_minutes
-        self.forecast_len_30 = forecast_minutes // interval_minutes
+        self.history_len = history_minutes // interval_minutes
+        self.forecast_len = forecast_minutes // interval_minutes
         # self.forecast_len_15 = forecast_minutes // 15
         # self.history_len_15 = history_minutes // 15
 
-        self.weighted_losses = WeightedLosses(forecast_length=self.forecast_len_30)
+        self.weighted_losses = WeightedLosses(forecast_length=self.forecast_len)
 
         self._accumulated_metrics = MetricAccumulator()
         self._accumulated_batches = BatchAccumulator(key_to_keep=self._target_key_name)
@@ -288,9 +293,9 @@ def use_quantile_regression(self):
     def num_output_features(self):
         """Number of ouput features he model chould predict for"""
         if self.use_quantile_regression:
-            out_features = self.forecast_len_30 * len(self.output_quantiles)
+            out_features = self.forecast_len * len(self.output_quantiles)
         else:
-            out_features = self.forecast_len_30
+            out_features = self.forecast_len
         return out_features
 
     def _quantiles_to_prediction(self, y_quantiles):
@@ -448,7 +453,7 @@ def training_step(self, batch, batch_idx):
         # Make all -1 values 0.0
         batch[self._target_key] = batch[self._target_key].clamp(min=0.0)
         y_hat = self(batch)
-        y = batch[self._target_key][:, -self.forecast_len_30 :, 0]
+        y = batch[self._target_key][:, -self.forecast_len:, 0]
 
         losses = self._calculate_common_losses(y, y_hat)
         losses = {f"{k}/train": v for k, v in losses.items()}
@@ -467,7 +472,7 @@ def validation_step(self, batch: dict, batch_idx):
         batch[self._target_key] = batch[self._target_key].clamp(min=0.0)
         y_hat = self(batch)
         # Sensor seems to be in batch, station, time order
-        y = batch[self._target_key][:, -self.forecast_len_30 :, 0]
+        y = batch[self._target_key][:, -self.forecast_len:, 0]
 
         losses = self._calculate_common_losses(y, y_hat)
         losses.update(self._calculate_val_losses(y, y_hat))
@@ -526,39 +531,24 @@ def validation_step(self, batch: dict, batch_idx):
                 )
                 plt.close(fig)
 
-                # Plot 1:30 to 3 hours ahead
-                fig = plot_batch_forecasts(
-                    batch,
-                    y_hat,
-                    quantiles=self.output_quantiles,
-                    key_to_plot=self._target_key_name,
-                    timesteps_to_plot=[6, 12],  # 1:30 to 3 hours ahead
-                )
-                self.logger.experiment.log(
-                    {
-                        f"val_forecast_samples/batch_idx_{accum_batch_num}_1.5_to_3hr": wandb.Image(
-                            fig
-                        ),
-                    }
-                )
-                plt.close(fig)
+                if self.time_step_intervals_to_plot is not None:
+                    for interval in self.time_step_intervals_to_plot:
+                        fig = plot_batch_forecasts(
+                            batch,
+                            y_hat,
+                            quantiles=self.output_quantiles,
+                            key_to_plot=self._target_key_name,
+                            timesteps_to_plot=interval,
+                        )
+                        self.logger.experiment.log(
+                            {
+                                f"val_forecast_samples/batch_idx_{accum_batch_num}_timestep_{interval}": wandb.Image(
+                                    fig
+                                ),
+                            }
+                        )
+                        plt.close(fig)
 
-                # Plot 15 to 39 hours ahead
-                fig = plot_batch_forecasts(
-                    batch,
-                    y_hat,
-                    quantiles=self.output_quantiles,
-                    key_to_plot=self._target_key_name,
-                    timesteps_to_plot=[60, 156],  # 15 to 39 hours ahead
-                )
-                self.logger.experiment.log(
-                    {
-                        f"val_forecast_samples/batch_idx_{accum_batch_num}_15_to_39hr": wandb.Image(
-                            fig
-                        ),
-                    }
-                )
-                plt.close(fig)
                 del self._val_y_hats
                 del self._val_batches
 
@@ -569,7 +559,7 @@ def test_step(self, batch, batch_idx):
         # Make all -1 values 0.0
         batch[self._target_key] = batch[self._target_key].clamp(min=0.0)
         y_hat = self(batch)
-        y = batch[self._target_key][:, -self.forecast_len_30 :, 0]
+        y = batch[self._target_key][:, -self.forecast_len:, 0]
 
         losses = self._calculate_common_losses(y, y_hat)
         losses.update(self._calculate_val_losses(y, y_hat))

pvnet/models/baseline/last_value.py

@@ -36,8 +36,8 @@ def forward(self, x: dict):
 
         # take the last value non forecaster value and the first in the pv yeild
         # (this is the pv site we are preditcting for)
-        y_hat = gsp_yield[:, -self.forecast_len_30 - 1, 0]
+        y_hat = gsp_yield[:, -self.forecast_len - 1, 0]
 
         # expand the last valid forward n predict steps
-        out = y_hat.unsqueeze(1).repeat(1, self.forecast_len_30)
+        out = y_hat.unsqueeze(1).repeat(1, self.forecast_len)
         return out

pvnet/models/baseline/single_value.py

@@ -33,5 +33,5 @@ def __init__(
     def forward(self, x: dict):
         """Run model forward on dict batch of data"""
         # Returns a single value at all steps
-        y_hat = torch.zeros_like(x[BatchKey.gsp][:, : self.forecast_len_30, 0]) + self._value
+        y_hat = torch.zeros_like(x[BatchKey.gsp][:, : self.forecast_len, 0]) + self._value
         return y_hat

pvnet/models/multimodal/deep_supervision.py

@@ -161,7 +161,7 @@ def __init__(
         if self.include_sun:
             # the minus 12 is bit of hard coded smudge for pvnet
             self.sun_fc1 = nn.Linear(
-                in_features=2 * (self.forecast_len_30 + self.history_len_30 + 1),
+                in_features=2 * (self.forecast_len + self.history_len + 1),
                 out_features=16,
             )
 
@@ -170,26 +170,26 @@ def __init__(
             num_cat_features += encoder_out_features
             self.sat_output_network = output_network(
                 in_features=encoder_out_features,
-                out_features=self.forecast_len_30,
+                out_features=self.forecast_len,
                 **output_network_kwargs,
             )
         if include_nwp:
             num_cat_features += encoder_out_features
             self.nwp_output_network = output_network(
                 in_features=encoder_out_features,
-                out_features=self.forecast_len_30,
+                out_features=self.forecast_len,
                 **output_network_kwargs,
             )
         if include_gsp_yield_history:
-            num_cat_features += self.history_len_30
+            num_cat_features += self.history_len
         if embedding_dim:
             num_cat_features += embedding_dim
         if include_sun:
             num_cat_features += 16
 
         self.output_network = output_network(
             in_features=num_cat_features,
-            out_features=self.forecast_len_30,
+            out_features=self.forecast_len,
             **output_network_kwargs,
         )
 
@@ -226,7 +226,7 @@ def encode(self, x):
         # *********************** GSP Data ************************************
         # add gsp yield history
         if self.include_gsp_yield_history:
-            gsp_history = x[BatchKey.gsp][:, : self.history_len_30].float()
+            gsp_history = x[BatchKey.gsp][:, : self.history_len].float()
             gsp_history = gsp_history.reshape(gsp_history.shape[0], -1)
             gsp_history = self.source_dropout_0d(gsp_history)
             modes["gsp"] = gsp_history
@@ -266,7 +266,7 @@ def multi_mode_forward(self, x):
     def training_step(self, batch, batch_idx):
         """Training step"""
         y_hats = self.multi_mode_forward(batch)
-        y = batch[BatchKey.gsp][:, -self.forecast_len_30 :, 0]
+        y = batch[BatchKey.gsp][:, -self.forecast_len:, 0]
 
         losses = self._calculate_common_losses(y, y_hats["all"])
         losses = {f"{k}/train": v for k, v in losses.items()}

pvnet/models/multimodal/multimodal.py

@@ -63,6 +63,7 @@ def __init__(
         pv_interval_minutes: int = 5,
         sat_interval_minutes: int = 5,
         sensor_interval_minutes: int = 30,
+        timestep_intervals_to_plot: Optional[list[int]] = None,
     ):
         """Neural network which combines information from different sources.
 
@@ -110,6 +111,7 @@ def __init__(
             pv_interval_minutes: The interval between each sample of the PV data
             sat_interval_minutes: The interval between each sample of the satellite data
             sensor_interval_minutes: The interval between each sample of the sensor data
+            timestep_intervals_to_plot: Intervals, in timesteps, to plot in addition to the full forecast
         """
 
         self.include_gsp_yield_history = include_gsp_yield_history
@@ -131,6 +133,7 @@ def __init__(
             output_quantiles=output_quantiles,
             target_key=target_key,
             interval_minutes=interval_minutes,
+            timestep_intervals_to_plot=timestep_intervals_to_plot
         )
 
         # Number of features expected by the output_network
@@ -234,7 +237,7 @@ def __init__(
         if self.include_sun:
             # the minus 12 is bit of hard coded smudge for pvnet
             self.sun_fc1 = nn.Linear(
-                in_features=2 * (self.forecast_len_30 + self.history_len_30 + 1),
+                in_features=2 * (self.forecast_len + self.history_len + 1),
                 out_features=16,
             )
 
@@ -243,7 +246,7 @@ def __init__(
 
         if include_gsp_yield_history:
             # Update num features
-            fusion_input_features += self.history_len_30
+            fusion_input_features += self.history_len
 
         self.output_network = output_network(
             in_features=fusion_input_features,
@@ -286,13 +289,13 @@ def forward(self, x):
                 # Target is PV, so only take the history
                 # Copy batch
                 x_tmp = x.copy()
-                x_tmp[BatchKey.pv] = x_tmp[BatchKey.pv][:, : self.history_len_30]
+                x_tmp[BatchKey.pv] = x_tmp[BatchKey.pv][:, : self.history_len]
                 modes["pv"] = self.pv_encoder(x_tmp)
 
         # *********************** GSP Data ************************************
         # add gsp yield history
         if self.include_gsp_yield_history:
-            gsp_history = x[BatchKey.gsp][:, : self.history_len_30].float()
+            gsp_history = x[BatchKey.gsp][:, : self.history_len].float()
             gsp_history = gsp_history.reshape(gsp_history.shape[0], -1)
             modes["gsp"] = gsp_history
 
@@ -314,7 +317,7 @@ def forward(self, x):
             else:
                 # Have to be its own Batch format
                 x_tmp = x.copy()
-                x_tmp[BatchKey.wind] = x_tmp[BatchKey.wind][:, : self.history_len_30]
+                x_tmp[BatchKey.wind] = x_tmp[BatchKey.wind][:, : self.history_len]
                 # This needs to be a Batch as input
                 modes["wind"] = self.wind_encoder(x_tmp)
 
@@ -324,7 +327,7 @@ def forward(self, x):
                 modes["sensor"] = self.sensor_encoder(x)
             else:
                 x_tmp = x.copy()
-                x_tmp[BatchKey.sensor] = x_tmp[BatchKey.sensor][:, : self.history_len_30]
+                x_tmp[BatchKey.sensor] = x_tmp[BatchKey.sensor][:, : self.history_len]
                 # This needs to be a Batch as input
                 modes["sensor"] = self.sensor_encoder(x_tmp)
 
@@ -339,6 +342,6 @@ def forward(self, x):
 
         if self.use_quantile_regression:
             # Shape: batch_size, seq_length * num_quantiles
-            out = out.reshape(out.shape[0], self.forecast_len_30, len(self.output_quantiles))
+            out = out.reshape(out.shape[0], self.forecast_len, len(self.output_quantiles))
 
         return out

pvnet/models/multimodal/nwp_weighting.py

@@ -63,7 +63,7 @@ def __init__(
         self.interpolate = nn.Sequential(
             nn.Linear(
                 in_features=nwp_sequence_len,
-                out_features=self.forecast_len_30,
+                out_features=self.forecast_len,
             ),
             nn.LeakyReLU(negative_slope=0.01),
         )

pvnet/models/multimodal/weather_residual.py

@@ -176,7 +176,7 @@ def __init__(
         if self.include_sun:
             # the minus 12 is bit of hard coded smudge for pvnet
             self.sun_fc1 = nn.Linear(
-                in_features=2 * (self.forecast_len_30 + self.history_len_30 + 1),
+                in_features=2 * (self.forecast_len + self.history_len + 1),
                 out_features=16,
             )
 
@@ -187,31 +187,31 @@ def __init__(
         if include_nwp:
             weather_cat_features += encoder_out_features
         if version == 1:
-            weather_cat_features += self.forecast_len_30
+            weather_cat_features += self.forecast_len
 
         nonweather_cat_features = 0
         if include_gsp_yield_history:
-            nonweather_cat_features += self.history_len_30
+            nonweather_cat_features += self.history_len
         if embedding_dim:
             nonweather_cat_features += embedding_dim
         if include_sun:
             nonweather_cat_features += 16
 
         self.simple_output_network = output_network(
             in_features=nonweather_cat_features,
-            out_features=self.forecast_len_30,
+            out_features=self.forecast_len,
             **output_network_kwargs,
         )
 
         self.weather_residual_network = nn.Sequential(
             output_network(
                 in_features=weather_cat_features,
-                out_features=self.forecast_len_30,
+                out_features=self.forecast_len,
                 **output_network_kwargs,
             ),
             # All output network return LeakyReLU activated outputs
             # However, the residual could be positive or negative
-            nn.Linear(self.forecast_len_30, self.forecast_len_30),
+            nn.Linear(self.forecast_len, self.forecast_len),
         )
 
         self.source_dropout_0d = CompleteDropoutNd(0, p=source_dropout)
@@ -247,7 +247,7 @@ def encode(self, x):
         # *********************** GSP Data ************************************
         # add gsp yield history
         if self.include_gsp_yield_history:
-            gsp_history = x[BatchKey.gsp][:, : self.history_len_30].float()
+            gsp_history = x[BatchKey.gsp][:, : self.history_len].float()
             gsp_history = gsp_history.reshape(gsp_history.shape[0], -1)
             gsp_history = self.source_dropout_0d(gsp_history)
             modes["gsp"] = gsp_history
@@ -295,7 +295,7 @@ def multi_mode_forward(self, x):
     def training_step(self, batch, batch_idx):
         """Run training step"""
         y_hats = self.multi_mode_forward(batch)
-        y = batch[BatchKey.gsp][:, -self.forecast_len_30 :, 0]
+        y = batch[BatchKey.gsp][:, -self.forecast_len:, 0]
 
         losses = self._calculate_common_losses(y, y_hats["weather_out"])
         losses = {f"{k}/train": v for k, v in losses.items()}