Rename localpatch.yaml to train_config.yaml and update notes; modify learning rate scheduler parameters and rename attention mask creation method

forecasting/models/vit_patch_model_local.py

@@ -51,9 +51,9 @@ class ViTLocal(pl.LightningModule):
 
         scheduler = CosineAnnealingWarmRestarts(
             optimizer,
-            T_0=250,  # Restart every 20 epochs
-            T_mult=2,  # Double the cycle length after each restart
-            eta_min=1e-7  # Minimum learning rate
+            T_0=150,
+            T_mult=2,
+            eta_min=1e-7
         )
 
         return {
@@ -66,8 +66,6 @@ class ViTLocal(pl.LightningModule):
             }
         }
 
-    # M/X Class Flare Detection Optimized Weights
-
     def _calculate_loss(self, batch, mode="train"):
         imgs, sxr = batch
         raw_preds, raw_patch_contributions = self.model(imgs, self.sxr_norm)
@@ -304,9 +302,9 @@ class InvertedAttentionBlock(nn.Module):
         )
 
         # Pre-compute attention mask for local interactions
-        self.register_buffer('attention_mask', self._create_local_attention_mask())
+        self.register_buffer('attention_mask', self._create_inverted_attention_mask())
 
-    def _create_local_attention_mask(self):
+    def _create_inverted_attention_mask(self):
         """Create attention mask for local interactions only"""
         # This creates a mask where only distant patches can attend to each other
 

forecasting/training/{localpatch.yaml → train_config.yaml}

@@ -52,4 +52,4 @@ wandb:
     - sxr
     - regression
   wb_name: paper-8-patch-4ch
-  notes: Regression from AIA images (4 channels) to GOES SXR flux
+  notes: Regression from AIA images to SXR images using ViTLocal model with 8x8 patches