tidy

Author: dfulu

pvnet/data/datamodule.py

@@ -12,24 +12,6 @@
 )
 
 
-def fill_nans_in_arrays(batch):
-    """Fills all NaN values in each np.ndarray in the batch dictionary with zeros.
-
-    Operation is performed in-place on the batch.
-    """
-    for k, v in batch.items():
-        if isinstance(v, torch.Tensor):
-            if torch.isnan(v).any():
-                batch[k] = torch.nan_to_num(v, nan=0.0)
-
-        # Recursion is included to reach NWP arrays in subdict
-        elif isinstance(v, dict):
-            fill_nans_in_arrays(v)
-
-    return batch
-
-
-
 class NumpybatchPremadeSamplesDataset(Dataset):
     """Dataset to load NumpyBatch samples"""
 
@@ -46,7 +28,7 @@ def __len__(self):
         return len(self.sample_paths)
 
     def __getitem__(self, idx):
-        return fill_nans_in_arrays(torch.load(self.sample_paths[idx]))
+        return torch.load(self.sample_paths[idx])
 
 
 def collate_fn(samples: list[NumpyBatch]):

pvnet/models/multimodal/multimodal.py

@@ -318,7 +318,7 @@ def forward(self, x):
             sat_data = torch.swapaxes(sat_data, 1, 2).float()  # switch time and channels
 
             if self.add_image_embedding_channel:
-                id = x[BatchKey[f"{self._target_key_name}_id"]][:, 0].int()
+                id = x[BatchKey[f"{self._target_key_name}_id"]].int()
                 sat_data = self.sat_embed(sat_data, id)
             modes["sat"] = self.sat_encoder(sat_data)
 
@@ -335,7 +335,7 @@ def forward(self, x):
                 nwp_data = torch.clip(nwp_data, min=-50, max=50)
 
                 if self.add_image_embedding_channel:
-                    id = x[BatchKey[f"{self._target_key_name}_id"]][:, 0].int()
+                    id = x[BatchKey[f"{self._target_key_name}_id"]].int()
                     nwp_data = self.nwp_embed_dict[nwp_source](nwp_data, id)
 
                 nwp_out = self.nwp_encoders_dict[nwp_source](nwp_data)
@@ -362,7 +362,7 @@ def forward(self, x):
 
         # ********************** Embedding of GSP ID ********************
         if self.embedding_dim:
-            id = x[BatchKey[f"{self._target_key_name}_id"]][:].int()
+            id = x[BatchKey[f"{self._target_key_name}_id"]].int()
             id_embedding = self.embed(id)
             modes["id"] = id_embedding
 

pvnet/utils.py

@@ -265,7 +265,7 @@ def _get_numpy(key):
     y_id_key = BatchKey[f"{key_to_plot}_id"]
     BatchKey[f"{key_to_plot}_t0_idx"]
     time_utc_key = BatchKey[f"{key_to_plot}_time_utc"]
-    y = batch[y_key][:, :].cpu().numpy()  # Select the one it is trained on
+    y = batch[y_key].cpu().numpy()  # Select the one it is trained on
     y_hat = y_hat.cpu().numpy()
     # Select between the timesteps in timesteps to plot
     plotting_name = key_to_plot.upper()