lint

Author: moritzhauschulz

src/weathergen/model/diffusion.py

@@ -23,17 +23,18 @@
 # ----------------------------------------------------------------------------
 
 
+import logging
 import math
 
 import torch
-import logging
 
 from weathergen.common.config import Config
 from weathergen.datasets.batch import SampleMetaData
 from weathergen.model.engines import ForecastingEngine
 
 logger = logging.getLogger(__name__)
 
+
 class DiffusionForecastEngine(torch.nn.Module):
     # Adopted from https://github.com/NVlabs/edm/blob/main/training/loss.py#L72
 
@@ -56,7 +57,7 @@ def __init__(self, cf: Config, num_healpix_cells: int, forecast_engine: Forecast
         self.rho = self.cf.rho
         self.p_mean = self.cf.p_mean
         self.p_std = self.cf.p_std
-        self.cur_token = None #for debugging only
+        self.cur_token = None  # for debugging only
 
     def forward(
         self, tokens: torch.Tensor, fstep: int, meta_info: dict[str, SampleMetaData]
@@ -75,12 +76,21 @@ def forward(
 
         c = 1  # TODO: add correct preconditioning (e.g., sample/s in previous time step)
 
-        #TODO: remove after single sample experiments
+        # TODO: remove after single sample experiments
         if self.cur_token is not None:
             logger.info("checking single sampling")
-            assert self.cur_token[0].shape == tokens[0].shape, 'first token shape was different between iterations – violates single sample overfitting with difference'
-            assert torch.equal(self.cur_token[0], tokens[0]), f'first token was different between iterations – violates single sample overfitting {self.cur_token[0] - tokens[0]}'
-            assert torch.equal(self.cur_token, tokens), 'tokens were different between iterations – violates single sample overfitting'
+            assert self.cur_token[0].shape == tokens[0].shape, (
+                "first token shape was different between iterations "
+                "– violates single sample overfitting with difference"
+            )
+            assert torch.equal(self.cur_token[0], tokens[0]), (
+                f"first token was different between iterations "
+                f"– violates single sample overfitting {self.cur_token[0] - tokens[0]}"
+            )
+            assert torch.equal(self.cur_token, tokens), (
+                f"tokens were different between iterations "
+                f"– violates single sample overfitting {self.cur_token - tokens}"
+            )
         self.cur_token = tokens
 
         y = tokens

src/weathergen/model/encoder.py

@@ -124,7 +124,6 @@ def forward(self, model_params, batch):
 
         return global_tokens, posteriors
 
-
     def assimilate_local(
         self, model_params, tokens: torch.Tensor, batch: ModelBatch
     ) -> torch.Tensor:

src/weathergen/model/engines.py

@@ -82,7 +82,6 @@ def __init__(self, cf: Config, sources_size, stream_names: list[str]) -> None:
                 )
             else:
                 raise ValueError("Unsupported embedding network type")
-        
 
     def forward(self, batch, pe_embed):
         num_steps_input = batch.get_num_steps()
@@ -108,7 +107,7 @@ def forward(self, batch, pe_embed):
 
             # embedding from physical space to per patch latent representation
             x_embeds += [self.embeds[stream_name](sdata).flatten(0, 1)]
-        
+
         # switch from stream to cell-based ordering and apply per cell positional encoding
 
         # computer scatter index across batch items and input steps

src/weathergen/model/model_interface.py

@@ -167,7 +167,7 @@ def init_model_and_shard(
     if cf.chkpt_encoder_weights:
         if is_root():
             logger.info(
-                f"Loading chkpt from run_id={cf.chkpt_encoder_weights}"\
+                f"Loading chkpt from run_id={cf.chkpt_encoder_weights}"
                 f" at mini_epoch {cf.chkpt_encoder_mini_epoch}."
             )
 
@@ -196,6 +196,7 @@ def init_model_and_shard(
 
     return model, model_params
 
+
 def load_encoder(cf, model, encoder_modules, device, run_id: str, mini_epoch=-1):
     """Loads model state from checkpoint and checks for missing and unused keys.
     Args:
@@ -282,7 +283,8 @@ def load_encoder(cf, model, encoder_modules, device, run_id: str, mini_epoch=-1)
         logger.warning(f"Unused keys when loading model: {mkeys}")
 
     return model
-    
+
+
 def load_model(cf, model, device, run_id: str, mini_epoch=-1):
     """Loads model state from checkpoint and checks for missing and unused keys.
     Args: