further linting of docstrings

Author: shmh40

src/weathergen/datasets/masking.py

@@ -7,8 +7,10 @@ class Masker:
     This class supports different masking strategies and combinations.
     Attributes:
         masking_rate (float): The base rate at which tokens are masked.
-        masking_strategy (str): The strategy used for masking (e.g., "random", "block", MORE TO BE IMPLEMENTED...).
-        TO BE IMPLEMENTED: masking_combination (str): The strategy for combining masking strategies through training (e.g., "sequential").
+        masking_strategy (str): The strategy used for masking (e.g., "random",
+        "block", MORE TO BE IMPLEMENTED...).
+        TO BE IMPLEMENTED: masking_combination (str): The strategy for combining masking
+        strategies through training (e.g., "sequential").
         masking_rate_sampling (bool): Whether to sample the masking rate from a distribution.
     """
 

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -339,7 +339,8 @@ def __iter__(self):
 
                             (ss_cells, ss_lens, ss_centroids) = self.tokenizer.batchify_source(
                                 stream_info,
-                                # NOTE: two unused arguments in TokenizerMasking, still used in TokenizerForecast?
+                                # NOTE: two unused arguments in TokenizerMasking,
+                                # still used in TokenizerForecast?
                                 self.masking_rate,
                                 self.masking_rate_sampling,
                                 torch.from_numpy(rdata.coords),

src/weathergen/datasets/tokenizer_masking.py

@@ -198,7 +198,9 @@ def batchify_source(
         if masking_rate == 1.0:
             token_lens = [len(t) for t in tokenized_data]
             self.perm_sel = [np.ones(l, dtype=bool) for l in token_lens]
-            source_tokens_cells = [c[~p] for c, p in zip(tokenized_data, self.perm_sel, strict=False)]
+            source_tokens_cells = [
+                c[~p] for c, p in zip(tokenized_data, self.perm_sel, strict=False)
+            ]
             source_tokens_lens = torch.zeros([self.num_healpix_cells_source], dtype=torch.int32)
             source_centroids = torch.tensor([])
             return (source_tokens_cells, source_tokens_lens, source_centroids)
@@ -281,10 +283,12 @@ def id(arg):
         )
 
         # --- MODIFICATION START ---
-        # The following block is modified to handle cases where a cell has no target tokens,
+        # The following block is modified to handle cases
+        # where a cell has no target tokens,
         # which would cause an error in torch.cat with an empty list.
 
-        # Pre-calculate the total feature dimension of a token to create correctly shaped empty tensors.
+        # Pre-calculate the total feature dimension of a token to create
+        # correctly shaped empty tensors.
         feature_dim = 6 + coords.shape[-1] + geoinfos.shape[-1] + source.shape[-1]
 
         processed_target_tokens = []