Make target key more general

Author: jacobbieker

pvnet/models/base_model.py

@@ -239,7 +239,7 @@ def __init__(
         forecast_minutes: int,
         optimizer: AbstractOptimizer,
         output_quantiles: Optional[list[float]] = None,
-        target_key: BatchKey = BatchKey.gsp,
+        target_key: str = "gsp",
     ):
         """Abtstract base class for PVNet submodels.
 
@@ -254,7 +254,8 @@ def __init__(
         super().__init__()
 
         self._optimizer = optimizer
-        self._target_key = target_key
+        self._target_key_name = target_key
+        self._target_key = BatchKey[f"{self._target_key}"]
 
         # Model must have lr to allow tuning
         # This setting is only used when lr is tuned with callback
@@ -275,7 +276,7 @@ def __init__(
 
         self._accumulated_metrics = MetricAccumulator()
         self._accumulated_batches = BatchAccumulator(
-            key_to_keep="gsp" if self._target_key == BatchKey.gsp else "sensor"
+            key_to_keep=self._target_key
         )
         self._accumulated_y_hat = PredAccumulator()
 
@@ -438,7 +439,7 @@ def _training_accumulate_log(self, batch, batch_idx, losses, y_hat):
                    y_hat,
                    batch_idx,
                    quantiles=self.output_quantiles,
-                   key_to_plot="gsp" if self._target_key == BatchKey.gsp else "wind",
+                   key_to_plot=self._target_key_name,
                )
                fig.savefig("latest_logged_train_batch.png")
 
@@ -503,7 +504,7 @@ def validation_step(self, batch: dict, batch_idx):
             if not hasattr(self, "_val_y_hats"):
                 self._val_y_hats = PredAccumulator()
                 self._val_batches = BatchAccumulator(
-                    key_to_keep="gsp" if self._target_key == BatchKey.gsp else "sensor"
+                    key_to_keep=self._target_key_name
                 )
 
             self._val_y_hats.append(y_hat)
@@ -517,7 +518,7 @@ def validation_step(self, batch: dict, batch_idx):
                    batch,
                    y_hat,
                    quantiles=self.output_quantiles,
-                   key_to_plot="gsp" if self._target_key == BatchKey.gsp else "sensor",
+                   key_to_plot=self._target_key_name,
                 )
 
                 self.logger.experiment.log(

pvnet/models/utils.py

@@ -100,14 +100,7 @@ def __bool__(self):
     # @staticmethod
     def _filter_batch_dict(self, d):
         keep_keys = (
-            [BatchKey.gsp, BatchKey.gsp_id, BatchKey.gsp_t0_idx, BatchKey.gsp_time_utc]
-            if self.key_to_keep == "gsp"
-            else [
-                BatchKey.wind,
-                BatchKey.wind_id,
-                BatchKey.wind_t0_idx,
-                BatchKey.wind_time_utc,
-            ]
+            [BatchKey[self.key_to_keep], BatchKey[f"{self.key_to_keep}_id"], BatchKey[f"{self.key_to_keep}_t0_idx"], BatchKey[f"{self.key_to_keep}_time_utc"]]
         )
         return {k: v for k, v in d.items() if k in keep_keys}
 
@@ -122,7 +115,7 @@ def flush(self) -> dict[BatchKey, list[torch.Tensor]]:
         """Concatenate all accumulated batches, return, and clear self"""
         batch = {}
         for k, v in self._batches.items():
-            if k == BatchKey.gsp_t0_idx or k == BatchKey.wind_t0_idx:
+            if k == BatchKey.gsp_t0_idx or k == BatchKey.wind_t0_idx or k == BatchKey.pv_t0_idx:
                 batch[k] = v[0]
             else:
                 batch[k] = torch.cat(v, dim=0)

pvnet/utils.py

@@ -248,15 +248,15 @@ def plot_batch_forecasts(batch, y_hat, batch_idx=None, quantiles=None, key_to_pl
     def _get_numpy(key):
         return batch[key].cpu().numpy().squeeze()
 
-    y_key = BatchKey.gsp if key_to_plot == "gsp" else BatchKey.wind
-    y_id_key = BatchKey.gsp_id if key_to_plot == "gsp" else BatchKey.wind_id
-    t0_idx_key = BatchKey.gsp_t0_idx if key_to_plot == "gsp" else BatchKey.wind_t0_idx
-    time_utc_key = BatchKey.gsp_time_utc if key_to_plot == "gsp" else BatchKey.wind_time_utc
+    y_key = BatchKey[f"{key_to_plot}"]
+    y_id_key = BatchKey[f"{key_to_plot}_id"]
+    t0_idx_key = BatchKey[f"{key_to_plot}_t0_idx"]
+    time_utc_key = BatchKey[f"{key_to_plot}_time_utc"]
     y = batch[y_key][:, :, 0].cpu().numpy()  # Select the one it is trained on
     y_hat = y_hat.cpu().numpy()
     gsp_ids = batch[y_id_key][:, 0].cpu().numpy().squeeze()
     t0_idx = int(batch[t0_idx_key])
-    plotting_name = "GSP" if key_to_plot == "gsp" else "Wind"
+    plotting_name = key_to_plot.upper()
 
     gsp_ids = batch[y_id_key].cpu().numpy().squeeze()
     batch[t0_idx_key]