Experimenting With Fusion Model

forecasting/inference/evaluation.py

@@ -941,10 +941,10 @@ class SolarFlareEvaluator:
 
 if __name__ == "__main__":
     # Example paths - replace with your actual paths
-    vit_csv = "/mnt/data/ML-READY/output/final_epoch_patch.csv"
+    vit_csv = "/mnt/data/ML-READY/output/patch.csv"
     baseline_results_csv = ""
     aia_data = "/mnt/data/ML-READY/AIA/test/"
-    weights_directory = "/mnt/data/ML-READY/final_epoch_patch_weights_final"
+    weights_directory = "/mnt/data/ML-READY/patch_weights"
 
     # Sample timestamps - Fixed the datetime generation
     start_time = datetime(2023, 8, 5, 20,30,00)

forecasting/inference/inference_on_patch.py

@@ -15,6 +15,7 @@ from torch.utils.data import DataLoader
 from forecasting.data_loaders.SDOAIA_dataloader import AIA_GOESDataset
 import forecasting.models as models
 from forecasting.models.vit_patch_model import ViT
+from forecasting.models import FusionViTHybrid
 from forecasting.models.linear_and_hybrid import HybridIrradianceModel  # Add your hybrid model import
 from forecasting.training.callback import unnormalize_sxr
 import yaml
@@ -28,7 +29,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def has_attention_weights(model):
     """Check if model supports attention weights"""
-    return hasattr(model, 'attention') or isinstance(model, ViT)
+    return hasattr(model, 'attention') or isinstance(model, ViT) or isinstance(model, FusionViTHybrid)
 
 
 def save_batch_flux_contributions(batch_flux_contributions, batch_idx, batch_size, times, flux_path, sxr_norm=None):
@@ -201,7 +202,7 @@ def load_model_from_config(config_data):
                 model_class = getattr(models, model_type)
                 model = model_class.load_from_checkpoint(checkpoint_path)
             except AttributeError:
-                raise ValueError(f"Unknown model type: {model_type}. Available types: ViT, HybridIrradianceModel")
+                raise ValueError(f"Unknown model type: {model_type}. Available types include: ViT, HybridIrradianceModel, FusionViTHybrid")
     else:
         # Regular PyTorch checkpoint
         state = torch.load(checkpoint_path, map_location=device, weights_only=False)
@@ -243,7 +244,7 @@ def main():
     parser.add_argument('-config', type=str, default='config.yaml', required=True, help='Path to config YAML.')
     parser.add_argument('-input_size', type=int, default=512, help='Input size for the model')
     parser.add_argument('-patch_size', type=int, default=16, help='Patch size for the model')
-    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for inference')
+    parser.add_argument('--batch_size', type=int, default=16, help='Batch size for inference')
     parser.add_argument('--no_weights', action='store_true', help='Skip saving attention weights to speed up')
     args = parser.parse_args()
 

forecasting/inference/inference_on_patch_config.yaml

@@ -1,6 +1,6 @@
 base_data_dir: "/mnt/data/ML-READY/"  # Change this line for different datasets
-output_path: "${base_data_dir}/output/final_epoch_patch.csv"
-weight_path: "${base_data_dir}/final_epoch_patch_weights_final/"
+output_path: "${base_data_dir}/output/patch.csv"
+weight_path: "${base_data_dir}/patch_weights/"
 flux_path: "${base_data_dir}/patch_flux/"
 mc:
   active: "false"
@@ -27,5 +27,5 @@ data:
   sxr_norm_path:
     "/mnt/data/ML-READY/SXR/normalized_sxr.npy"
   checkpoint_path:
-    "/mnt/data/ML-READY/new-checkpoint/vit-16-higher-weight-lower-decay-epoch=288-val_total_loss=0.0385.ckpt"
+    "/mnt/data/ML-READY/new-checkpoint/vit-16-MSE-deeper-epoch=51-val_total_loss=0.1064.ckpt"
 

forecasting/models/__init__.py

@@ -0,0 +1,2 @@
+from .fusion_vit_hybrid import FusionViTHybrid
+

forecasting/models/fusion_vit_hybrid.py

@@ -0,0 +1,228 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+from .vit_patch_model import VisionTransformer, SXRRegressionDynamicLoss, normalize_sxr, unnormalize_sxr
+from .linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel
+
+
+class FusionViTHybrid(pl.LightningModule):
+    """End-to-end fused model: ViT for spatial patches + Linear/Hybrid for scalar.
+
+    - ViT branch outputs per-patch raw flux and a ViT global (sum of patches).
+    - Scalar branch (Linear or Hybrid) outputs a global scalar.
+    - A learnable gate blends the two globals; the spatial map uses ViT's
+      distribution but is calibrated to the fused/global prediction.
+
+    Forward returns a 4-tuple compatible with existing inference utils:
+        (global_fused, attention_weights, fused_patch_flux, global_fused)
+    """
+
+    def __init__(
+        self,
+        vit_kwargs: dict,
+        scalar_branch: str,
+        scalar_kwargs: dict,
+        sxr_norm,
+        lr: float = 1e-4,
+        lambda_vit_to_target: float = 0.3,
+        lambda_scalar_to_target: float = 0.1,
+        use_attention: bool = True,
+        learnable_gate: bool = True,
+        gate_init_bias: float = 5.0,
+        weight_decay: float = 1e-5,
+        cosine_restart_T0: int = 50,
+        cosine_restart_Tmult: int = 2,
+        cosine_eta_min: float = 1e-7,
+    ):
+        super().__init__()
+
+        # Save hyperparameters needed for checkpointing
+        self.save_hyperparameters(ignore=["sxr_norm"])  # sxr_norm is a tensor/array
+
+        # Branches: filter unsupported keys for VisionTransformer
+        filtered_vit_kwargs = dict(vit_kwargs)
+        filtered_vit_kwargs.pop('lr', None)
+        filtered_vit_kwargs.pop('num_classes', None)
+        self.vit = VisionTransformer(**filtered_vit_kwargs)
+
+        if scalar_branch.lower() in ["linear", "lineairradiancemodel"]:
+            self.scalar = LinearIrradianceModel(
+                d_input=scalar_kwargs.get("d_input"),
+                d_output=scalar_kwargs.get("d_output"),
+                loss_func=scalar_kwargs.get("loss_func", nn.HuberLoss()),
+                lr=scalar_kwargs.get("lr", lr),
+            )
+        elif scalar_branch.lower() in ["hybrid", "hybridirradiancemodel"]:
+            self.scalar = HybridIrradianceModel(
+                d_input=scalar_kwargs.get("d_input"),
+                d_output=scalar_kwargs.get("d_output"),
+                cnn_model=scalar_kwargs.get("cnn_model", "updated"),
+                ln_model=scalar_kwargs.get("ln_model", True),
+                ln_params=scalar_kwargs.get("ln_params", None),
+                lr=scalar_kwargs.get("lr", lr),
+                cnn_dp=scalar_kwargs.get("cnn_dp", 0.75),
+                loss_func=scalar_kwargs.get("loss_func", nn.HuberLoss()),
+            )
+        else:
+            raise ValueError(f"Unknown scalar_branch: {scalar_branch}")
+
+        # Loss and normalization
+        self.sxr_norm = sxr_norm
+        self.adaptive_loss = SXRRegressionDynamicLoss(window_size=1500)
+
+        # Gate: learnable scalar in [0,1] blending scalar vs vit global
+        self.learnable_gate = learnable_gate
+        if learnable_gate:
+            self.gate_logit = nn.Parameter(torch.tensor(gate_init_bias, dtype=torch.float32))
+        else:
+            self.register_buffer("gate_logit", torch.tensor(gate_init_bias, dtype=torch.float32))
+
+        # Optim params
+        self.lr = lr
+        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
+
+        # Aux loss weights
+        self.lambda_vit_to_target = lambda_vit_to_target
+        self.lambda_scalar_to_target = lambda_scalar_to_target
+
+        # Whether to compute/return attention
+        self.use_attention = use_attention
+
+    def forward(self, x, return_attention: bool = True):
+        # ViT branch: returns different numbers of values based on return_attention
+        vit_out = self.vit(x, self.sxr_norm, return_attention=(self.use_attention and return_attention))
+
+        if self.use_attention and return_attention and len(vit_out) == 3:
+            global_vit_raw, attention_weights, patch_flux_raw = vit_out
+        else:
+            global_vit_raw, patch_flux_raw = vit_out
+            attention_weights = None
+
+        # Scalar branch expects (B,H,W,C)
+        global_scalar_raw = self.scalar(x)
+        # Ensure positivity for SXR-like targets
+        global_scalar_raw = F.softplus(global_scalar_raw)
+
+        # Shapes: ensure tensors are shaped [B, 1]
+        if global_vit_raw.dim() == 1:
+            global_vit_raw = global_vit_raw.unsqueeze(-1)
+        if global_scalar_raw.dim() == 1:
+            global_scalar_raw = global_scalar_raw.unsqueeze(-1)
+
+        # Patch weights from ViT distribution
+        weights = patch_flux_raw / (global_vit_raw.clamp(min=1e-15))
+
+        # Blend globals via sigmoid(gate_logit)
+        gate = torch.sigmoid(self.gate_logit)
+        global_fused = gate * global_scalar_raw + (1.0 - gate) * global_vit_raw
+        # Avoid zeros/negatives before log normalization downstream
+        global_fused = global_fused.clamp(min=1e-15)
+
+        # Calibrated patch flux using fused global
+        fused_patch_flux = global_fused * weights
+
+        # Match inference API: (pred, attn, patch_flux, total_from_patches)
+        return global_fused, attention_weights, fused_patch_flux, global_fused
+    def forward_for_callback(self, x, return_attention: bool = True):
+        """Forward method compatible with AttentionMapCallback"""
+        global_fused, attention_weights, fused_patch_flux, _ = self.forward(x, return_attention)
+        # Callback expects (outputs, attention_weights, _)
+        return attention_weights
+    def _calc_losses(self, imgs, sxr):
+        # Forward
+        global_fused, attention_weights, fused_patch_flux, _ = self(imgs, return_attention=True)
+
+        # Main adaptive loss on fused global
+        raw_preds_squeezed = torch.squeeze(global_fused)
+        sxr_un = unnormalize_sxr(sxr, self.sxr_norm)
+        norm_preds_squeezed = normalize_sxr(raw_preds_squeezed, self.sxr_norm)
+        main_loss, weights_adapt = self.adaptive_loss.calculate_loss(
+            norm_preds_squeezed, sxr, sxr_un, raw_preds_squeezed
+        )
+
+        # Auxiliary consistency losses (vit and scalar heads individually)
+        # Recompute heads without extra forward
+        # Extract vit global by re-running vit without attention to save memory
+        with torch.no_grad():
+            vit_out = self.vit(imgs, self.sxr_norm, return_attention=False)
+        global_vit_raw = vit_out[0]
+        if global_vit_raw.dim() > 1:
+            global_vit_raw = torch.squeeze(global_vit_raw)
+        global_vit_raw = global_vit_raw.clamp(min=1e-15)
+        vit_norm = normalize_sxr(global_vit_raw, self.sxr_norm)
+        loss_vit = F.huber_loss(vit_norm, sxr)
+
+        global_scalar_raw = self.scalar(imgs)
+        global_scalar_raw = F.softplus(global_scalar_raw)
+        if global_scalar_raw.dim() > 1:
+            global_scalar_raw = torch.squeeze(global_scalar_raw)
+        global_scalar_raw = global_scalar_raw.clamp(min=1e-15)
+        scalar_norm = normalize_sxr(global_scalar_raw, self.sxr_norm)
+        loss_scalar = F.huber_loss(scalar_norm, sxr)
+
+        total_loss = main_loss \
+            + self.lambda_vit_to_target * loss_vit \
+            + self.lambda_scalar_to_target * loss_scalar
+
+        return total_loss, {
+            "main_loss": main_loss.detach(),
+            "loss_vit": loss_vit.detach(),
+            "loss_scalar": loss_scalar.detach(),
+        }
+
+    def training_step(self, batch, batch_idx):
+        imgs, sxr = batch
+        total_loss, logs = self._calc_losses(imgs, sxr)
+
+        # Logs
+        self.log("train_main_loss", logs["main_loss"], on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        self.log("train_vit_loss", logs["loss_vit"], on_step=True, on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
+        self.log("train_scalar_loss", logs["loss_scalar"], on_step=True, on_epoch=True, prog_bar=False, logger=True, sync_dist=True)
+
+        # Learning rate
+        current_lr = self.trainer.optimizers[0].param_groups[0]['lr']
+        self.log('learning_rate', current_lr, on_step=True, on_epoch=False, prog_bar=True, logger=True, sync_dist=True)
+
+        return total_loss
+
+    def validation_step(self, batch, batch_idx):
+        imgs, sxr = batch
+        total_loss, logs = self._calc_losses(imgs, sxr)
+        self.log("val_main_loss", logs["main_loss"], on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        self.log("val_total_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        return total_loss
+
+    def test_step(self, batch, batch_idx):
+        imgs, sxr = batch
+        total_loss, _ = self._calc_losses(imgs, sxr)
+        self.log("test_total_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        return total_loss
+
+    def configure_optimizers(self):
+        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'
+            }
+        }
+
+

forecasting/models/vit_patch_model.py

@@ -306,8 +306,8 @@ class SXRRegressionDynamicLoss:
             'x_class': 20.0
         }
     def calculate_loss(self, preds_squeezed, sxr, sxr_un, preds_squeezed_un):
-        #base_loss = F.huber_loss(preds_squeezed, sxr, delta=1.0, reduction='none')
-        base_loss = F.mse_loss(preds_squeezed, sxr, reduction='none')
+        base_loss = F.huber_loss(preds_squeezed, sxr, delta=1.0, reduction='none')
+        #base_loss = F.mse_loss(preds_squeezed, sxr, reduction='none')
         weights = self._get_adaptive_weights(sxr_un, preds_squeezed_un, base_loss)
         self._update_tracking(sxr_un, preds_squeezed_un, base_loss)
         weighted_loss = base_loss * weights

forecasting/training/callback.py

@@ -124,7 +124,11 @@ class AttentionMapCallback(Callback):
             imgs = imgs[:self.num_samples].to(pl_module.device)
 
             # Get predictions with attention weights
-            outputs, attention_weights, _  = pl_module(imgs, return_attention=True)
+            #Dynamically extract attention weights from the model
+            try:
+                outputs, attention_weights, _  = pl_module(imgs, return_attention=True)
+            except:
+                attention_weights = pl_module.forward_for_callback(imgs, return_attention=True)
 
             # Visualize attention for each sample
             for sample_idx in range(min(self.num_samples, imgs.size(0))):
@@ -134,7 +138,7 @@ class AttentionMapCallback(Callback):
                     attention_weights,
                     sample_idx,
                     trainer.current_epoch,
-                    pl_module.model.patch_size
+                    patch_size=16
                 )
                 trainer.logger.experiment.log({"Attention plots": wandb.Image(map)})
                 plt.close(map)

forecasting/training/config.yaml

@@ -4,8 +4,8 @@ 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: "ViT Patch"  # Options: "cnn", "vit",
-batch_size:    80
+selected_model: "FusionViTHybrid"  # Options: "cnn", "vit",
+batch_size:    16
 epochs:        500
 oversample: false
 balance_strategy: "upsample_minority"
@@ -23,12 +23,27 @@ vit_custom:
     num_classes: 1
     patch_size: 16
     num_patches: 1024
-    hidden_dim: 1024
+    hidden_dim: 512
     num_heads: 8
     num_layers: 6
     dropout: 0.1
     lr: 0.0001
 
+
+fusion:
+  scalar_branch: "hybrid"        # or "linear"
+  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
+
+
 # Data paths (automatically constructed from base directories)
 data:
   aia_dir:
@@ -48,5 +63,5 @@ wandb:
     - aia
     - sxr
     - regression
-  wb_name: vit-16-MSE-deeper
+  wb_name: vit-fused-model
   notes: Regression from AIA images (6 channels) to GOES SXR flux

forecasting/training/train.py

@@ -22,6 +22,7 @@ from forecasting.data_loaders.SDOAIA_dataloader import AIA_GOESDataModule
 from forecasting.models.vision_transformer_custom import ViT
 from forecasting.models.linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel
 from forecasting.models.vit_patch_model import ViT as ViTPatch
+from forecasting.models import FusionViTHybrid
 from callback import ImagePredictionLogger_SXR, AttentionMapCallback
 from pytorch_lightning.callbacks import Callback
 
@@ -204,14 +205,39 @@ elif config_data['selected_model'] == 'hybrid':
 elif config_data['selected_model'] == 'ViT':
     model = ViT(model_kwargs=config_data['vit_custom'], sxr_norm = sxr_norm)
 
-elif config_data['selected_model'] == 'ViT Patch':
+elif config_data['selected_model'] == 'ViTPatch':
     model = ViTPatch(model_kwargs=config_data['vit_custom'], sxr_norm = sxr_norm)
 
+elif config_data['selected_model'] == 'FusionViTHybrid':
+    # Expect a 'fusion' section in YAML
+    fusion_cfg = config_data.get('fusion', {})
+    scalar_branch = fusion_cfg.get('scalar_branch', 'hybrid')
+    scalar_kwargs = fusion_cfg.get('scalar_kwargs', {
+        'd_input': len(config_data['wavelengths']),
+        'd_output': 1,
+        'cnn_model': config_data.get('megsai', {}).get('cnn_model', 'updated'),
+        'cnn_dp': config_data.get('megsai', {}).get('cnn_dp', 0.75),
+        'lr': fusion_cfg.get('lr', config_data.get('megsai', {}).get('lr', 1e-4)),
+    })
+    vit_kwargs = config_data.get('vit_custom', {})
+
+    model = FusionViTHybrid(
+        vit_kwargs=vit_kwargs,
+        scalar_branch=scalar_branch,
+        scalar_kwargs=scalar_kwargs,
+        sxr_norm=sxr_norm,
+        lr=fusion_cfg.get('lr', 1e-4),
+        lambda_vit_to_target=fusion_cfg.get('lambda_vit_to_target', 0.3),
+        lambda_scalar_to_target=fusion_cfg.get('lambda_scalar_to_target', 0.1),
+        learnable_gate=fusion_cfg.get('learnable_gate', True),
+        gate_init_bias=fusion_cfg.get('gate_init_bias', 5.0),
+    )
+
 else:
     raise NotImplementedError(f"Architecture {config_data['selected_model']} not supported.")
 
 # Trainer
-if config_data['selected_model'] == 'ViT' or config_data['selected_model'] == 'ViT Patch':
+if config_data['selected_model'] == 'ViT' or config_data['selected_model'] == 'ViT Patch' or config_data['selected_model'] == 'FusionViTHybrid':
     trainer = Trainer(
         default_root_dir=config_data['data']['checkpoints_dir'],
         accelerator="gpu" if torch.cuda.is_available() else "cpu",