Fix formatting with pre-commit

mllam · Mar 17, 2024 · 2d9afd1 · 2d9afd1
1 parent 65f00ea
commit 2d9afd1
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 14 deletions.
diff --git a/neural_lam/constants.py b/neural_lam/constants.py
@@ -116,5 +116,5 @@
 )
 
 # Data dimensions
-GRID_FORCING_DIM = 5 * 3 + 1 # 5 feat. for 3 time-step window + 1 batch-static
+GRID_FORCING_DIM = 5 * 3 + 1  # 5 feat. for 3 time-step window + 1 batch-static
 GRID_STATE_DIM = 17
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -116,9 +116,7 @@ def expand_to_batch(x, batch_size):
         """
         return x.unsqueeze(0).expand(batch_size, -1, -1)
 
-    def predict_step(
-        self, prev_state, prev_prev_state, forcing
-    ):
+    def predict_step(self, prev_state, prev_prev_state, forcing):
         """
         Step state one step ahead using prediction model, X_{t-1}, X_t -> X_t+1
         prev_state: (B, num_grid_nodes, feature_dim), X_t
@@ -127,9 +125,7 @@ def predict_step(
         """
         raise NotImplementedError("No prediction step implemented")
 
-    def unroll_prediction(
-        self, init_states, forcing_features, true_states
-    ):
+    def unroll_prediction(self, init_states, forcing_features, true_states):
         """
         Roll out prediction taking multiple autoregressive steps with model
         init_states: (B, 2, num_grid_nodes, d_f)

diff --git a/neural_lam/models/base_graph_model.py b/neural_lam/models/base_graph_model.py
@@ -98,9 +98,7 @@ def process_step(self, mesh_rep):
         """
         raise NotImplementedError("process_step not implemented")
 
-    def predict_step(
-        self, prev_state, prev_prev_state, forcing
-    ):
+    def predict_step(self, prev_state, prev_prev_state, forcing):
         """
         Step state one step ahead using prediction model, X_{t-1}, X_t -> X_t+1
         prev_state: (B, num_grid_nodes, feature_dim), X_t

diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
@@ -171,11 +171,9 @@ def __getitem__(self, idx):
         )  # (dim_x, dim_y, 1)
         # Flatten
         water_cover_features = water_cover_features.flatten(0, 1)  # (N_grid, 1)
-        # Exand over temporal dimension
+        # Expand over temporal dimension
         water_cover_expanded = water_cover_features.unsqueeze(0).expand(
-            self.sample_length - 2,  # -2 as added on after windowing
-            -1,
-            -1
+            self.sample_length - 2, -1, -1  # -2 as added on after windowing
         )  # (sample_len, N_grid, 1)
 
         # TOA flux