Updated splitting function so now it can do SXR and AIA individually... it can also repartition an already partioned folder... updated the megs models to include warm restarts, weight decay, and adaptive weights

data/split_data.py

@@ -1,72 +1,183 @@
 import os
 import pandas as pd
 import shutil
+import argparse
 from datetime import datetime
+from pathlib import Path
 
-aia_data_dir = "/mnt/data2/AIA_processed/"
-sxr_data_dir = "/mnt/data2/ML-Ready_clean/GOES-18-SXR-B/"
-flares_event_dir = "/mnt/data2/ML-Ready_clean/flares_event_dir/"
-non_flares_event_dir = "/mnt/data2/ML-Ready_clean/non_flares_event_dir/"
-mixed_data_dir = "/mnt/data2/ML-Ready_clean/mixed_data/"
-flare_events_csv = "/mnt/data2/SDO-AIA-flaring/FlareEvents/flare_events_2012-01-01_2015-03-25.csv"
-
-# Create train, val, test subdirectories under flaring and non-flaring
-for base_dir in [flares_event_dir, non_flares_event_dir, mixed_data_dir]:
-    os.makedirs(os.path.join(base_dir, "AIA"), exist_ok=True)
-    os.makedirs(os.path.join(base_dir, "SXR"), exist_ok=True)
+def split_data(input_folder, output_dir, data_type, flare_events_csv=None, repartition=False):
+    """
+    Split data from input folder into train/val/test based on month.
+    Optionally use flare events for additional classification.
+    
+    Args:
+        input_folder (str): Path to the input folder containing data files
+        output_dir (str): Path to the output directory where split data will be saved
+        data_type (str): Type of data ('aia' or 'sxr')
+        flare_events_csv (str, optional): Path to the flare events CSV file
+        repartition (bool): If True, treat input_folder as already partitioned (has train/val/test subdirs)
+    """
+    
+    # Validate input folder
+    if not os.path.exists(input_folder):
+        raise ValueError(f"Input folder does not exist: {input_folder}")
+    
+    # Validate data type
+    if data_type.lower() not in ['aia', 'sxr']:
+        raise ValueError("data_type must be 'aia' or 'sxr'")
+    
+    # Create output directory structure
+    os.makedirs(output_dir, exist_ok=True)
     for split in ["train", "val", "test"]:
-        os.makedirs(os.path.join(base_dir, "AIA", split), exist_ok=True)
-        os.makedirs(os.path.join(base_dir, "SXR", split), exist_ok=True)
-
-# Load flare events
-flare_event = pd.read_csv(flare_events_csv)
-
-# Create list of flare event time ranges
-flaring_eve_list = []
-for i, row in flare_event.iterrows():
-    start_time = pd.to_datetime(row['event_starttime'])
-    end_time = pd.to_datetime(row['event_endtime'])
-    flaring_eve_list.append((start_time, end_time))
-
-# Get list of files in data_dir
-data_list = os.listdir(aia_data_dir)
-
-for file in data_list:
-    try:
-        aia_time = pd.to_datetime(file.split(".")[0])
-    except ValueError:
-        print(f"Skipping file {file}: Invalid timestamp format")
-        continue
-
-    # Determine if the file is during a flare event
-    is_flaring = any(start <= aia_time <= end for start, end in flaring_eve_list)
-    base_dir = flares_event_dir if is_flaring else non_flares_event_dir
-
-    month = aia_time.month
-
-    if month in [2, 3, 4, 5, 6, 7, 9, 10, 11, 12]:
-        split_dir = "train"
-    elif month == 1:
-        split_dir = "val"
-    elif month == 8:
-        split_dir = "test"
-    else:
-        print(f"Skipping file {file}: Unexpected month {month}")
-        continue
-
-    # Copy file to appropriate directory
-    src_aia = os.path.join(aia_data_dir, file)
-    src_sxr = os.path.join(sxr_data_dir, file)
-    dst_aia = os.path.join(base_dir, "AIA", split_dir, file)
-    dst_sxr = os.path.join(base_dir, "SXR", split_dir, file)
-
-    if not os.path.exists(dst_aia):
-        shutil.copy(src_aia, dst_aia)
-        print(f"Copied {file} to {dst_aia} and {dst_sxr}")
+        os.makedirs(os.path.join(output_dir, split), exist_ok=True)
+    
+    print(f"Processing {data_type.upper()} data from: {input_folder}")
+    print(f"Output directory: {output_dir}")
+    
+    # Load flare events if provided
+    flaring_eve_list = []
+    if flare_events_csv and os.path.exists(flare_events_csv):
+        print(f"Loading flare events from: {flare_events_csv}")
+        flare_event = pd.read_csv(flare_events_csv)
+        
+        # Create list of flare event time ranges
+        for i, row in flare_event.iterrows():
+            start_time = pd.to_datetime(row['event_starttime'])
+            end_time = pd.to_datetime(row['event_endtime'])
+            flaring_eve_list.append((start_time, end_time))
+        print(f"Loaded {len(flaring_eve_list)} flare events")
     else:
-        print(f"File {dst_aia} already exists, skipping copy.")
-    if not os.path.exists(dst_sxr):
-        shutil.copy(src_sxr, dst_sxr)
+        print("No flare events CSV provided or file not found. Skipping flare classification.")
+    
+    # Get list of files in input folder
+    if repartition:
+        # For repartitioning, collect files from train/val/test subdirectories
+        data_list = []
+        for split in ["train", "val", "test"]:
+            split_dir = os.path.join(input_folder, split)
+            if os.path.exists(split_dir):
+                split_files = os.listdir(split_dir)
+                # Add split information to each file for tracking
+                for file in split_files:
+                    data_list.append((file, split))
+                print(f"Found {len(split_files)} files in {split}/ directory")
+        print(f"Total files to repartition: {len(data_list)}")
     else:
-        print(f"File {dst_sxr} already exists, skipping copy.")
+        # For normal splitting, get files directly from input folder
+        data_list = os.listdir(input_folder)
+        print(f"Found {len(data_list)} files to process")
+    
+    moved_count = 0
+    skipped_count = 0
+    
+    for file_info in data_list:
+        if repartition:
+            file, original_split = file_info
+        else:
+            file = file_info
+            original_split = None
+            
+        try:
+            # Extract timestamp from filename (assuming format like "2012-01-01T00:00:00.npy")
+            file_time = pd.to_datetime(file.split(".")[0])
+        except ValueError:
+            print(f"Skipping file {file}: Invalid timestamp format")
+            skipped_count += 1
+            continue
+        
+        # Determine if the file is during a flare event (if flare events are available)
+        is_flaring = False
+        if flaring_eve_list:
+            is_flaring = any(start <= file_time <= end for start, end in flaring_eve_list)
+        
+        # Determine split based on month
+        month = file_time.month
+        
+        if month in [2, 3, 4, 5, 6, 7, 9, 10, 11, 12]:
+            new_split_dir = "train"
+        elif month == 1:
+            new_split_dir = "val"
+        elif month == 8:
+            new_split_dir = "test"
+        else:
+            print(f"Skipping file {file}: Unexpected month {month}")
+            skipped_count += 1
+            continue
+        
+        # Determine source and destination paths
+        if repartition:
+            src_path = os.path.join(input_folder, original_split, file)
+        else:
+            src_path = os.path.join(input_folder, file)
+        
+        dst_path = os.path.join(output_dir, new_split_dir, file)
+        
+        # Skip if file is already in the correct split and we're repartitioning
+        if repartition and original_split == new_split_dir and os.path.exists(dst_path):
+            print(f"File {file} already in correct split ({new_split_dir}), skipping.")
+            skipped_count += 1
+            continue
+        
+        if not os.path.exists(dst_path):
+            try:
+                shutil.move(src_path, dst_path)
+                if repartition:
+                    if flaring_eve_list:
+                        print(f"Moved {file} from {original_split}/ to {new_split_dir}/ (flaring: {is_flaring})")
+                    else:
+                        print(f"Moved {file} from {original_split}/ to {new_split_dir}/")
+                else:
+                    if flaring_eve_list:
+                        print(f"Moved {file} to {new_split_dir}/ (flaring: {is_flaring})")
+                    else:
+                        print(f"Moved {file} to {new_split_dir}/")
+                moved_count += 1
+            except Exception as e:
+                print(f"Error moving {file}: {e}")
+                skipped_count += 1
+        else:
+            print(f"File {dst_path} already exists, skipping move.")
+            skipped_count += 1
+    
+    print(f"\nProcessing complete!")
+    print(f"Files moved: {moved_count}")
+    print(f"Files skipped: {skipped_count}")
+    print(f"Total files processed: {moved_count + skipped_count}")
+
+def main():
+    parser = argparse.ArgumentParser(description='Split AIA or SXR data into train/val/test sets based on month')
+    parser.add_argument('--input_folder', type=str, required=True,
+                       help='Path to the input folder containing data files (or partitioned folder for repartition)')
+    parser.add_argument('--output_dir', type=str, required=True,
+                       help='Path to the output directory where split data will be saved')
+    parser.add_argument('--data_type', type=str, choices=['aia', 'sxr'], required=True,
+                       help='Type of data: "aia" or "sxr"')
+    parser.add_argument('--flare_events_csv', type=str, default=None,
+                       help='Path to the flare events CSV file (optional)')
+    parser.add_argument('--repartition', action='store_true',
+                       help='Repartition an already partitioned folder (input_folder should have train/val/test subdirs)')
+    
+    args = parser.parse_args()
+    
+    # Convert to absolute paths
+    input_folder = os.path.abspath(args.input_folder)
+    output_dir = os.path.abspath(args.output_dir)
+    flare_events_csv = os.path.abspath(args.flare_events_csv) if args.flare_events_csv else None
+    
+    # Validate repartition mode
+    if args.repartition:
+        # Check if input folder has train/val/test subdirectories
+        expected_dirs = ['train', 'val', 'test']
+        missing_dirs = []
+        for dir_name in expected_dirs:
+            if not os.path.exists(os.path.join(input_folder, dir_name)):
+                missing_dirs.append(dir_name)
+        
+        if missing_dirs:
+            print(f"Warning: Input folder is missing expected subdirectories: {missing_dirs}")
+            print("Continuing with available directories...")
+    
+    split_data(input_folder, output_dir, args.data_type, flare_events_csv, args.repartition)
 
+if __name__ == "__main__":
+    main()

forecasting/data_loaders/sxr_normalization.py

@@ -51,7 +51,7 @@ def compute_sxr_norm(sxr_dir):
 
 if __name__ == "__main__":
     # Update this path to your real data SXR directory
-    sxr_dir = "/mnt/data/ML-Ready-mixed/ML-Ready-mixed/SXR/train"  # Replace with actual path
+    sxr_dir = "/mnt/data/ML-READY/SXR/train"  # Replace with actual path
     sxr_norm = compute_sxr_norm(sxr_dir)
-    np.save("/mnt/data/ML-Ready-mixed/ML-Ready-mixed/SXR/normalized_sxr.npy", sxr_norm)
+    np.save("/mnt/data/ML-READY/SXR/normalized_sxr.npy", sxr_norm)
     #print(f"Saved SXR normalization to /mnt/data/ML-Ready-Data-No-Intensity-Cut/normalized_sxr")

forecasting/inference/evaluation.py

@@ -890,7 +890,7 @@ class SolarFlareEvaluator:
         # Sort frame paths by timestamp to ensure correct order
         frame_paths.sort(key=lambda x: os.path.basename(x))
 
-        movie_path = os.path.join(self.output_dir, "AIA_video_with_uncertainties.mp4")
+        movie_path = os.path.join(self.output_dir, f"AIA_{timestamps[0].split('T')[0]}.mp4")
         with imageio.get_writer(movie_path, fps=30, codec='libx264', format='ffmpeg') as writer:
             for frame_path in frame_paths:
                 if os.path.exists(frame_path):

forecasting/models/base_model.py

@@ -1,30 +1,72 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from pytorch_lightning import LightningModule
+from collections import deque
+import numpy as np
+
+# Import adaptive loss and normalization functions
+from .vit_patch_model import SXRRegressionDynamicLoss, normalize_sxr, unnormalize_sxr
 
 class BaseModel(LightningModule):
-    def __init__(self, model, loss_func, lr):
+    def __init__(self, model, loss_func, lr, sxr_norm=None, weight_decay=1e-5, 
+                 cosine_restart_T0=50, cosine_restart_Tmult=2, cosine_eta_min=1e-7):
         super().__init__()
         self.model = model
         self.loss_func = loss_func
         self.lr = lr
+        self.sxr_norm = sxr_norm
+        self.weight_decay = weight_decay
+        self.cosine_restart_T0 = cosine_restart_T0
+        self.cosine_restart_Tmult = cosine_restart_Tmult
+        self.cosine_eta_min = cosine_eta_min
+        
+        # Initialize adaptive loss if sxr_norm is provided
+        if sxr_norm is not None:
+            self.adaptive_loss = SXRRegressionDynamicLoss(window_size=1500)
+        else:
+            self.adaptive_loss = None
 
     def forward(self, x):
         return self.model(x)
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
-        scheduler = {
-            'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3),
-            'monitor': 'val_loss',  # name of the metric to monitor
-            'interval': 'epoch',
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+        )
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer,
+            T_0=self.cosine_restart_T0,
+            T_mult=self.cosine_restart_Tmult,
+            eta_min=self.cosine_eta_min,
+        )
+        return {
+            'optimizer': optimizer,
+            'lr_scheduler': {
+                'scheduler': scheduler,
+                'interval': 'epoch',
+                'frequency': 1,
+                'name': 'learning_rate'
+            }
         }
-        return {'optimizer': optimizer, 'lr_scheduler': scheduler}
 
     def training_step(self, batch, batch_idx):
         x, target = batch
         pred = self(x)
-        loss = self.loss_func(torch.squeeze(pred), target)
+        
+        # Use adaptive loss if available and sxr_norm is provided
+        if self.adaptive_loss is not None and self.sxr_norm is not None:
+            raw_preds_squeezed = torch.squeeze(pred)
+            target_un = unnormalize_sxr(target, self.sxr_norm)
+            norm_preds_squeezed = normalize_sxr(raw_preds_squeezed, self.sxr_norm)
+            loss, weights = self.adaptive_loss.calculate_loss(
+                norm_preds_squeezed, target, target_un, raw_preds_squeezed
+            )
+        else:
+            loss = self.loss_func(torch.squeeze(pred), target)
+        
         self.log('train_loss', loss)
         current_lr = self.trainer.optimizers[0].param_groups[0]['lr']
         self.log('learning_rate', current_lr, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
@@ -33,13 +75,35 @@ class BaseModel(LightningModule):
     def validation_step(self, batch, batch_idx):
         x, target = batch
         pred = self(x)
-        loss = self.loss_func(torch.squeeze(pred), target)
+        
+        # Use adaptive loss if available and sxr_norm is provided
+        if self.adaptive_loss is not None and self.sxr_norm is not None:
+            raw_preds_squeezed = torch.squeeze(pred)
+            target_un = unnormalize_sxr(target, self.sxr_norm)
+            norm_preds_squeezed = normalize_sxr(raw_preds_squeezed, self.sxr_norm)
+            loss, weights = self.adaptive_loss.calculate_loss(
+                norm_preds_squeezed, target, target_un, raw_preds_squeezed
+            )
+        else:
+            loss = self.loss_func(torch.squeeze(pred), target)
+        
         self.log('val_loss', loss)
         return loss
 
     def test_step(self, batch, batch_idx):
         x, target = batch
         pred = self(x)
-        loss = self.loss_func(torch.squeeze(pred), target)
+        
+        # Use adaptive loss if available and sxr_norm is provided
+        if self.adaptive_loss is not None and self.sxr_norm is not None:
+            raw_preds_squeezed = torch.squeeze(pred)
+            target_un = unnormalize_sxr(target, self.sxr_norm)
+            norm_preds_squeezed = normalize_sxr(raw_preds_squeezed, self.sxr_norm)
+            loss, weights = self.adaptive_loss.calculate_loss(
+                norm_preds_squeezed, target, target_un, raw_preds_squeezed
+            )
+        else:
+            loss = self.loss_func(torch.squeeze(pred), target)
+        
         self.log('test_loss', loss)
         return loss

forecasting/models/linear_and_hybrid.py

@@ -10,11 +10,14 @@ from forecasting.models.base_model import BaseModel
 from torchvision.models import resnet18
 
 class LinearIrradianceModel(BaseModel):
-    def __init__(self, d_input, d_output, loss_func=HuberLoss(), lr=1e-4):
+    def __init__(self, d_input, d_output, loss_func=HuberLoss(), lr=1e-4, sxr_norm=None, 
+                 weight_decay=1e-5, cosine_restart_T0=50, cosine_restart_Tmult=2, cosine_eta_min=1e-7):
         self.n_channels = d_input
         self.outSize = d_output
         model = nn.Linear(2 * self.n_channels, self.outSize)
-        super().__init__(model=model, loss_func=loss_func, lr=lr)
+        super().__init__(model=model, loss_func=loss_func, lr=lr, sxr_norm=sxr_norm,
+                        weight_decay=weight_decay, cosine_restart_T0=cosine_restart_T0,
+                        cosine_restart_Tmult=cosine_restart_Tmult, cosine_eta_min=cosine_eta_min)
 
     def forward(self, x, **kwargs):
 
@@ -48,14 +51,21 @@ class LinearIrradianceModel(BaseModel):
         return self.model(input_features)
 
 class HybridIrradianceModel(BaseModel):
-    def __init__(self, d_input, d_output, cnn_model='resnet', ln_model=True, ln_params=None, lr=1e-4, cnn_dp=0.75, loss_func=HuberLoss()):
-        super().__init__(model=None, loss_func=loss_func, lr=lr)
+    def __init__(self, d_input, d_output, cnn_model='resnet', ln_model=True, ln_params=None, lr=1e-4, cnn_dp=0.75, loss_func=HuberLoss(),
+                 sxr_norm=None, weight_decay=1e-5, cosine_restart_T0=50, cosine_restart_Tmult=2, cosine_eta_min=1e-7):
+        super().__init__(model=None, loss_func=loss_func, lr=lr, sxr_norm=sxr_norm,
+                        weight_decay=weight_decay, cosine_restart_T0=cosine_restart_T0,
+                        cosine_restart_Tmult=cosine_restart_Tmult, cosine_eta_min=cosine_eta_min)
         self.n_channels = d_input
         self.outSize = d_output
         self.ln_params = ln_params
         self.ln_model = None
         if ln_model:
-            self.ln_model = LinearIrradianceModel(d_input, d_output, loss_func=loss_func, lr=lr)
+            self.ln_model = LinearIrradianceModel(d_input, d_output, loss_func=loss_func, lr=lr,
+                                                sxr_norm=sxr_norm, weight_decay=weight_decay,
+                                                cosine_restart_T0=cosine_restart_T0,
+                                                cosine_restart_Tmult=cosine_restart_Tmult,
+                                                cosine_eta_min=cosine_eta_min)
         if self.ln_params is not None and self.ln_model is not None:
             self.ln_model.model.weight = nn.Parameter(self.ln_params['weight'])
             self.ln_model.model.bias = nn.Parameter(self.ln_params['bias'])

forecasting/training/config.yaml

@@ -4,9 +4,9 @@ base_data_dir: "/mnt/data/ML-READY"  # Change this line for different datasets
 base_checkpoint_dir: "/mnt/data/ML-READY"    # Change this line for different datasets
 wavelengths: [94, 131, 171, 193, 211, 304]  # AIA wavelengths in Angstroms
 # Model configuration
-selected_model: "FusionViTHybrid"  # Options: "cnn", "vit",
-batch_size:    16
-epochs:        500
+selected_model: "cnn"  # Options: "cnn", "vit",
+batch_size:    256
+epochs:        250
 oversample: false
 balance_strategy: "upsample_minority"
 
@@ -14,8 +14,12 @@ megsai:
   architecture: "cnn"
   seed: 42
   lr: 0.0001
-  cnn_model: "original"
+  cnn_model: "updated"
   cnn_dp: 0.2
+  weight_decay: 1e-5
+  cosine_restart_T0: 50
+  cosine_restart_Tmult: 2
+  cosine_eta_min: 1e-7
 
 vit_custom:
     embed_dim: 512
@@ -57,11 +61,11 @@ data:
 
 wandb:
   entity: jayantbiradar619-university-of-arizona # Use your exact W&B username
-  project: ViT Loss Function Testing For Patches
+  project: ViT Proper Scale
   job_type: training
   tags:
     - aia
     - sxr
     - regression
-  wb_name: vit-fused-model
+  wb_name: baseline-model
   notes: Regression from AIA images (6 channels) to GOES SXR flux

forecasting/training/train.py

@@ -190,8 +190,13 @@ if config_data['selected_model'] == 'linear':
     model = LinearIrradianceModel(
         d_input= len(config_data['wavelengths']),
         d_output=1,
-        lr= config_data['model']['lr'],
-        loss_func=MSELoss()
+        lr= config_data['megsai']['lr'],
+        loss_func=MSELoss(),
+        sxr_norm=sxr_norm,
+        weight_decay=config_data['megsai']['weight_decay'],
+        cosine_restart_T0=config_data['megsai']['cosine_restart_T0'],
+        cosine_restart_Tmult=config_data['megsai']['cosine_restart_Tmult'],
+        cosine_eta_min=config_data['megsai']['cosine_eta_min']
     )
 elif config_data['selected_model'] == 'hybrid':
     model = HybridIrradianceModel(
@@ -201,6 +206,11 @@ elif config_data['selected_model'] == 'hybrid':
         ln_model=True,
         cnn_dp=config_data['megsai']['cnn_dp'],
         lr=config_data['megsai']['lr'],
+        sxr_norm=sxr_norm,
+        weight_decay=config_data['megsai']['weight_decay'],
+        cosine_restart_T0=config_data['megsai']['cosine_restart_T0'],
+        cosine_restart_Tmult=config_data['megsai']['cosine_restart_Tmult'],
+        cosine_eta_min=config_data['megsai']['cosine_eta_min']
     )
 elif config_data['selected_model'] == 'ViT':
     model = ViT(model_kwargs=config_data['vit_custom'], sxr_norm = sxr_norm)