updates to evaluation script and callbacks

forecasting/inference/auto_evaluate.py

@@ -305,7 +305,7 @@ def main():
     parser.add_argument('-checkpoint_dir', type=str, help='Directory containing checkpoint files')
     parser.add_argument('-checkpoint_path', type=str, help='Specific checkpoint file path')
     parser.add_argument('-model_name', type=str, required=True, help='Name for the model (used for output naming)')
-    parser.add_argument('-base_data_dir', type=str, default='/mnt/data/NO-OVERLAP', help='Base data directory')
+    parser.add_argument('-base_data_dir', type=str, default='/mnt/data/', help='Base data directory')
     parser.add_argument('-skip_inference', action='store_true', help='Skip inference and only run evaluation')
     parser.add_argument('-skip_evaluation', action='store_true', help='Skip evaluation and only run inference')
     

forecasting/inference/inference_template.yaml

@@ -57,25 +57,3 @@ data:
   sxr_norm_path: "${base_data_dir}/SXR-SPLIT/normalized_sxr.npy"
   checkpoint_path: "PLACEHOLDER_CHECKPOINT_PATH"  # Will be replaced by batch script
 
-# MEGSAI parameters (should match training config)
-megsai:
-  cnn_model: "updated"
-  cnn_dp: 0.2
-  weight_decay: 1e-5
-  cosine_restart_T0: 50
-  cosine_restart_Tmult: 2
-  cosine_eta_min: 1e-7
-
-# Fusion parameters (if using fusion model)
-fusion:
-  scalar_branch: "hybrid"
-  lr: 0.0001
-  lambda_vit_to_target: 0.3
-  lambda_scalar_to_target: 0.1
-  learnable_gate: true
-  gate_init_bias: 5.0
-  scalar_kwargs:
-    d_input: 6
-    d_output: 1
-    cnn_model: "updated"
-    cnn_dp: 0.75

forecasting/training/callback.py

@@ -274,9 +274,57 @@ class AttentionMapCallback(Callback):
             img_display = (img_np[:, :, 0] + 1) / 2
             img_display = np.stack([img_display] * 3, axis=2)
 
-        # Visualization layout logic (unchanged)
-        # [The plotting logic remains as-is from the original script]
-        # Produces multiple subplots showing attention patterns and overlayed maps.
+        # Create the figure and subplots
+        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+        fig.suptitle(f'Attention Visualization - Epoch {epoch}, Sample {sample_idx}', fontsize=16)
+        
+        # Plot 1: Original image
+        axes[0, 0].imshow(img_display)
+        axes[0, 0].set_title('Original Image')
+        axes[0, 0].axis('off')
+        
+        # Plot 2: Attention map
+        im1 = axes[0, 1].imshow(attention_map, cmap='hot', interpolation='nearest')
+        axes[0, 1].set_title('Attention Map')
+        axes[0, 1].axis('off')
+        plt.colorbar(im1, ax=axes[0, 1])
+        
+        # Plot 3: Overlay
+        axes[0, 2].imshow(img_display)
+        axes[0, 2].imshow(attention_map, cmap='hot', alpha=0.6, interpolation='nearest')
+        axes[0, 2].set_title('Attention Overlay')
+        axes[0, 2].axis('off')
+        
+        # Plot 4: Center attention (if available)
+        if center_map is not None:
+            im2 = axes[1, 0].imshow(center_map, cmap='hot', interpolation='nearest')
+            axes[1, 0].set_title('Center Patch Attention')
+            axes[1, 0].axis('off')
+            plt.colorbar(im2, ax=axes[1, 0])
+        else:
+            axes[1, 0].text(0.5, 0.5, 'Center attention\nnot available', 
+                           ha='center', va='center', transform=axes[1, 0].transAxes)
+            axes[1, 0].set_title('Center Patch Attention')
+            axes[1, 0].axis('off')
+        
+        # Plot 5: Patch flux (if available)
+        if patch_flux is not None:
+            patch_flux_np = patch_flux.cpu().numpy().reshape(grid_h, grid_w)
+            im3 = axes[1, 1].imshow(patch_flux_np, cmap='viridis', interpolation='nearest')
+            axes[1, 1].set_title('Patch Flux')
+            axes[1, 1].axis('off')
+            plt.colorbar(im3, ax=axes[1, 1])
+        else:
+            axes[1, 1].text(0.5, 0.5, 'Patch flux\nnot available', 
+                           ha='center', va='center', transform=axes[1, 1].transAxes)
+            axes[1, 1].set_title('Patch Flux')
+            axes[1, 1].axis('off')
+        
+        # Plot 6: Attention statistics
+        axes[1, 2].hist(attention_map.flatten(), bins=50, alpha=0.7)
+        axes[1, 2].set_title('Attention Distribution')
+        axes[1, 2].set_xlabel('Attention Weight')
+        axes[1, 2].set_ylabel('Frequency')
 
         plt.tight_layout()
         return fig

forecasting/training/localpatch.yaml

@@ -1,7 +1,7 @@
 
 #Base directories - change these to switch datasets
-base_data_dir: "/mnt/data/PAPER_DATA_B"  # Change this line for different datasets
-base_checkpoint_dir: "/mnt/data/PAPER_DATA_B"    # Change this line for different datasets
+base_data_dir: "/mnt/data/PAPER_DATA_WITH_335"  # Change this line for different datasets
+base_checkpoint_dir: "/mnt/data/PAPER_DATA_WITH_335"    # Change this line for different datasets
 wavelengths: [94, 131, 171, 193, 211, 304, 335]  # AIA wavelengths in Angstroms
 
 # GPU configuration
@@ -13,8 +13,8 @@ wavelengths: [94, 131, 171, 193, 211, 304, 335]  # AIA wavelengths in Angstroms
 gpu_ids: "all"  # Use both GPUs
 # Model configuration
 selected_model: "ViTLocal"  # Options: "hybrid", "vit", "fusion", "vitpatch"
-batch_size:    4
-epochs:        250
+batch_size:    6
+epochs:        150
 oversample: false
 balance_strategy: "upsample_minority"
 calculate_base_weights: false  # Whether to calculate class-based weights for loss function
@@ -51,5 +51,5 @@ wandb:
     - aia
     - sxr
     - regression
-  wb_name: paper-testing-8-patch-335-512-hidden-6-layers-256-embed-dim
+  wb_name: paper-testing-8-patch-335-512-hidden-6-layers-256-embed-dim-updated
   notes: Regression from AIA images (6 channels) to GOES SXR flux