bug fixes

flaring/MEGS_AI_baseline/callback.py

@@ -10,12 +10,12 @@ import astropy.units as u
 # Custom Callback
 sdoaia94 = matplotlib.colormaps['sdoaia94']
 
-def unnormalize(y, eve_norm):
-    eve_norm = torch.tensor(eve_norm).float()
-    norm_mean = eve_norm[0]
-    norm_stdev = eve_norm[1]
-    y = y * norm_stdev[None].to(y) + norm_mean[None].to(y)
-    return y
+
+def unnormalize_sxr(normalized_values, sxr_norm):
+    if isinstance(normalized_values, torch.Tensor):
+        normalized_values = normalized_values.cpu().numpy()
+    normalized_values = np.array(normalized_values, dtype=np.float32)
+    return 10 ** (normalized_values * float(sxr_norm[1].item()) + float(sxr_norm[0].item())) - 1e-8
 
 
 class ImagePredictionLogger_SXR(Callback):
@@ -27,12 +27,6 @@ class ImagePredictionLogger_SXR(Callback):
         self.val_sxr = data_samples[1]
         self.sxr_norm = sxr_norm
 
-    def unnormalize_sxr(self, normalized_values):
-        if isinstance(normalized_values, torch.Tensor):
-            normalized_values = normalized_values.cpu().numpy()
-        normalized_values = np.array(normalized_values, dtype=np.float32)
-        return 10 ** (normalized_values * float(self.sxr_norm[1].item()) + float(self.sxr_norm[0].item())) - 1e-8
-
     def on_validation_epoch_end(self, trainer, pl_module):
 
         aia_images = []
@@ -50,8 +44,8 @@ class ImagePredictionLogger_SXR(Callback):
             aia_images.append(aia.squeeze(0).cpu().numpy())
             true_sxr.append(target.item())
 
-        true_unorm = self.unnormalize_sxr(true_sxr)
-        pred_unnorm = self.unnormalize_sxr(pred_sxr)
+        true_unorm = unnormalize_sxr(true_sxr,self.sxr_norm)
+        pred_unnorm = unnormalize_sxr(pred_sxr,self.sxr_norm)
         fig1 = self.plot_aia_sxr(aia_images,true_unorm, pred_unnorm)
         trainer.logger.experiment.log({"Soft X-ray flux plots": wandb.Image(fig1)})
         plt.close(fig1)
@@ -85,178 +79,3 @@ class ImagePredictionLogger_SXR(Callback):
 
         fig.tight_layout()
         return fig
-
-
-class ImagePredictionLogger(Callback):
-    def __init__(self, val_imgs, val_eve, names, aia_wavelengths):
-        super().__init__()
-        self.val_imgs, self.val_eve = val_imgs, val_eve
-        self.names = names
-        self.aia_wavelengths = aia_wavelengths
-
-    def on_validation_epoch_end(self, trainer, pl_module):
-        # Bring the tensors to CPU
-        val_imgs = self.val_imgs.to(device=pl_module.device)
-        # Get model prediction
-        # pred_eve = pl_module.forward(val_imgs).cpu().numpy()
-        pred_eve = pl_module.forward_unnormalize(val_imgs).cpu().numpy()
-        val_eve = unnormalize(self.val_eve, pl_module.eve_norm).numpy()
-        val_imgs = val_imgs.cpu().numpy()
-
-        # create matplotlib figure
-        fig = self.plot_aia_eve(val_imgs, val_eve, pred_eve)
-        # Log the images to wandb
-        trainer.logger.experiment.log({"AIA Images and EVE bar plots": wandb.Image(fig)})
-        plt.close(fig)
-
-    def plot_aia_eve(self, val_imgs, val_eve, pred_eve):
-        """
-        Function to plot a 4 channel AIA stack and the EVE barplots
-
-        Arguments:
-        ----------
-            val_imgs: numpy array
-                Stack with 4 image channels
-            val_eve: numpy array
-                Stack of ground-truth eve channels
-            pred_eve: numpy array
-                Stack of predicted eve channels
-        Returns:
-        --------
-            fig: matplotlib figure
-                figure with plots
-        """
-        samples = pred_eve.shape[0]
-        n_aia_wavelengths = len(self.aia_wavelengths)
-        wspace = 0.2
-        hspace = 0.125
-        dpi = 100
-
-        if n_aia_wavelengths < 3:
-            nrows = 1
-            ncols = n_aia_wavelengths
-            fig = plt.figure(figsize=(9 + 9 / 4 * n_aia_wavelengths, 3 * samples), dpi=dpi)
-            gs = fig.add_gridspec(samples, n_aia_wavelengths + 3, wspace=wspace, hspace=hspace)
-        elif n_aia_wavelengths < 5:
-            nrows = 2
-            ncols = 2
-            fig = plt.figure(figsize=(9 + 9 / 4 * 2, 6 * samples), dpi=dpi)
-            gs = fig.add_gridspec(2 * samples, 5, wspace=wspace, hspace=hspace)
-        elif n_aia_wavelengths < 7:
-            nrows = 2
-            ncols = 3
-            fig = plt.figure(figsize=(9 + 9 / 4 * 3, 6 * samples), dpi=dpi)
-            gs = fig.add_gridspec(2 * samples, 6, wspace=wspace, hspace=hspace)
-        else:
-            nrows = 2
-            ncols = 4
-            fig = plt.figure(figsize=(15, 5 * samples), dpi=dpi)
-            gs = fig.add_gridspec(2 * samples, 7, wspace=wspace, hspace=hspace)
-
-        cmaps_all = ['sdoaia94', 'sdoaia131', 'sdoaia171', 'sdoaia193', 'sdoaia211',
-                     'sdoaia304', 'sdoaia335', 'sdoaia1600', 'sdoaia1700']
-        cmaps = [cmaps_all[i] for i in self.aia_wavelengths]
-        n_plots = 0
-
-        for s in range(samples):
-            for i in range(nrows):
-                for j in range(ncols):
-                    if n_plots < n_aia_wavelengths:
-                        ax = fig.add_subplot(gs[s * nrows + i, j])
-                        ax.imshow(val_imgs[s, i * ncols + j], cmap=plt.get_cmap(cmaps[i * ncols + j]), origin='lower')
-                        ax.text(0.01, 0.99, cmaps[i * ncols + j], horizontalalignment='left', verticalalignment='top',
-                                color='w', transform=ax.transAxes)
-                        ax.set_axis_off()
-                        n_plots += 1
-            n_plots = 0
-            # eve data
-            ax5 = fig.add_subplot(gs[s * nrows, ncols:])
-            if self.names is not None:
-                ax5.bar(np.arange(0, len(val_eve[s, :])), val_eve[s, :], label='ground truth')
-                ax5.bar(np.arange(0, len(pred_eve[s, :])), pred_eve[s, :], width=0.5, label='prediction', alpha=0.5)
-                ax5.set_xticks(np.arange(0, len(val_eve[s, :])))
-                ax5.set_xticklabels(self.names, rotation=45)
-            else:
-                ax5.plot(np.arange(0, len(val_eve[s, :])), val_eve[s, :], label='ground truth', alpha=0.5,
-                         drawstyle='steps-mid')
-                ax5.plot(np.arange(0, len(pred_eve[s, :])), pred_eve[s, :], label='prediction', alpha=0.5,
-                         drawstyle='steps-mid')
-            ax5.set_yscale('log')
-            ax5.legend()
-
-            ax6 = fig.add_subplot(gs[s * nrows + 1, ncols:])
-            if self.names is not None:
-                ax6.bar(np.arange(0, len(val_eve[s, :])), np.abs(pred_eve[s, :] - val_eve[s, :]) / val_eve[s, :] * 100,
-                        label='relative error (%)')
-                ax6.set_xticks(np.arange(0, len(val_eve[s, :])))
-                ax6.set_xticklabels(self.names, rotation=45)
-            else:
-                ax6.plot(np.arange(0, len(val_eve[s, :])), np.abs(pred_eve[s, :] - val_eve[s, :]) / val_eve[s, :] * 100,
-                         label='relative error (%)', alpha=0.5, drawstyle='steps-mid')
-            ax6.set_yscale('log')
-            ax6.legend()
-
-        fig.tight_layout()
-        return fig
-
-
-class SpectrumPredictionLogger(ImagePredictionLogger):
-    def __init__(self, val_imgs, val_eve, names, aia_wavelengths):
-        super().__init__(val_imgs, val_eve, names, aia_wavelengths)
-
-    def plot_aia_eve(self, val_imgs, val_eve, pred_eve):
-        """
-        Function to plot a 4 channel AIA stack and the EVE barplots
-
-        Arguments:
-        ----------
-            val_imgs: numpy array
-                Stack with 4 image channels
-            val_eve: numpy array
-                Stack of ground-truth eve channels
-            pred_eve: numpy array
-                Stack of predicted eve channels
-        Returns:
-        --------
-            fig: matplotlib figure
-                figure with plots
-        """
-        samples = pred_eve.shape[0]
-        n_aia_wavelengths = len(self.aia_wavelengths)
-        wspace = 0.2
-        hspace = 0.125
-        dpi = 200
-
-        fig = plt.figure(figsize=(5, 5), dpi=dpi)
-        gs = fig.add_gridspec(2, 1, wspace=wspace, hspace=hspace)
-
-        # eve data
-        s = 0
-        ax5 = fig.add_subplot(gs[0, 0])
-        if self.names is not None:
-            ax5.bar(np.arange(0, len(val_eve[s, :])), val_eve[s, :], label='ground truth')
-            ax5.bar(np.arange(0, len(pred_eve[s, :])), pred_eve[s, :], width=0.5, label='prediction', alpha=0.5)
-            ax5.set_xticks(np.arange(0, len(val_eve[s, :])))
-            ax5.set_xticklabels(self.names, rotation=45)
-        else:
-            ax5.plot(np.arange(0, len(val_eve[s, :])), val_eve[s, :], label='ground truth', alpha=0.5,
-                     drawstyle='steps-mid')
-            ax5.plot(np.arange(0, len(pred_eve[s, :])), pred_eve[s, :], label='prediction', alpha=0.5,
-                     drawstyle='steps-mid')
-        ax5.set_yscale('log')
-        ax5.legend()
-
-        ax6 = fig.add_subplot(gs[1, 0])
-        if self.names is not None:
-            ax6.bar(np.arange(0, len(val_eve[s, :])), np.abs(pred_eve[s, :] - val_eve[s, :]) / val_eve[s, :] * 100,
-                    label='relative error (%)')
-            ax6.set_xticks(np.arange(0, len(val_eve[s, :])))
-            ax6.set_xticklabels(self.names, rotation=45)
-        else:
-            ax6.plot(np.arange(0, len(val_eve[s, :])), np.abs(pred_eve[s, :] - val_eve[s, :]) / val_eve[s, :] * 100,
-                     label='relative error (%)', alpha=0.5, drawstyle='steps-mid')
-        ax6.set_yscale('log')
-        ax6.legend()
-
-        fig.tight_layout()
-        return fig

flaring/MEGS_AI_baseline/config.yaml

@@ -11,7 +11,9 @@
     cnn_dp:
       - 0.75
     epochs:
-      - 100
+      - 1
+    save_dictionary:
+      -
   wandb:
     entity: jayantbiradar619-university-of-arizona # Use your exact W&B username
     project: MEGS-AI flaring  # Lowercase, no spaces

flaring/MEGS_AI_baseline/inference.py

@@ -3,28 +3,41 @@ import torch
 import numpy as np
 from torch.utils.data import DataLoader
 from SDOAIA_dataloader import AIA_GOESDataset
+from models.linear_and_hybrid import HybridIrradianceModel
+from callback import ImagePredictionLogger_SXR
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 def predict_log_outputs(model, dataset, batch_size=8):
-    """Generator yielding raw log-space model outputs"""
     model.eval()
     loader = DataLoader(dataset, batch_size=batch_size)
 
+    # Get device from model
+    device = next(model.parameters()).device
+
     with torch.no_grad():
         for batch in loader:
-            # Handle different dataset formats
+            # Correct unpacking based on your data structure
             if isinstance(batch, tuple) and len(batch) == 2:
-                aia_imgs = batch[0][0]  # Unpack ((aia, sxr), target)
+                # batch = (inputs, targets) where inputs = [aia_imgs, sxr_imgs]
+                aia_imgs = batch[0][0]  # Get aia_imgs from inputs
             else:
-                aia_imgs = batch[0] if isinstance(batch, (list, tuple)) else batch
+                # Fallback for other formats
+                aia_imgs = batch[0][0] if isinstance(batch[0], list) else batch[0]
+
+            # Move to device (it's already a tensor)
+            aia_imgs = aia_imgs.to(device)
+
+            # Get model predictions
+            log_outputs = model(aia_imgs)
 
-            aia_imgs = aia_imgs.to(next(model.parameters()).device)
-            log_outputs = model(aia_imgs)  # Get raw log-space outputs
+            # Move to CPU and convert to numpy before yielding
             yield from log_outputs.cpu().numpy()
 
 def main():
     parser = argparse.ArgumentParser(description='Save raw log-space model outputs')
-    parser.add_argument('--model', required=True, help='Path to trained model')
-    parser.add_argument('--aia-dir', required=True, help='Directory of AIA images')
+    parser.add_argument('--ckpt_path', required=True, help='Path to model checkpoint')
+    parser.add_argument('--aia_dir', required=True, help='Directory of AIA images')
+    parser.add_argument('--sxr_dir', required=True, help='Directory of target SXR images')
     parser.add_argument('--output', default='log_predictions.txt',
                         help='Output file for log-space predictions')
     parser.add_argument('--batch-size', type=int, default=8,
@@ -33,14 +46,24 @@ def main():
     args = parser.parse_args()
 
     # Setup
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = torch.load(args.model, map_location=device).to(device)
+    state = torch.load(args.ckpt_path, map_location=device, weights_only=False)
+    model = state['model']
+    model.to(device)
+    # Assume it's a checkpoint with state_dict
+
+    # model = HybridIrradianceModel(6,1)
+    # state_dict = checkpoint.get('state_dict', checkpoint)
+    #
+    # # Handle potential key mismatches (e.g., PyTorch Lightning prefixes)
+    # state_dict = {k.replace('model.', ''): v for k, v in state_dict.items()}
+    # model.load_state_dict(state_dict, strict=False)
+
+
 
     # Dataset without any output transformation
     dataset = AIA_GOESDataset(
         aia_dir=args.aia_dir,
-        sxr_dir='',  # No SXR files needed
-        sxr_norm=None,  # Skip any normalization
+        sxr_dir=args.sxr_dir,  # No SXR files needed
         transform=None  # No input transforms
     )
 
@@ -48,7 +71,7 @@ def main():
     with open(args.output, 'w') as f:
         f.write("# Log-space SXR predictions (log10(W/m²))\n")
         for log_pred in predict_log_outputs(model, dataset, args.batch_size):
-            f.write(f"{log_pred:.6f}\n")  # Write with 6 decimal places
+            print(log_pred)
 
     print(f"Log-space predictions saved to {args.output}")
     print("These are raw model outputs in log10 space before any exponentiation")

flaring/MEGS_AI_baseline/models/linear_and_hybrid.py

@@ -108,9 +108,6 @@ class HybridIrradianceModel(BaseModel):
         if isinstance(x, (list, tuple)):
             x = x[0]
 
-        # Debug: Print input shape
-        print(f"Input shape to HybridIrradianceModel.forward: {x.shape}")
-
         # Expect x shape: (batch_size, H, W, C)
         if len(x.shape) != 4:
             raise ValueError(f"Expected 4D input tensor (batch_size, H, W, C), got shape {x.shape}")

flaring/MEGS_AI_baseline/train.py

@@ -1,6 +1,8 @@
 
 import argparse
 import os
+from datetime import datetime
+
 import yaml
 import itertools
 import wandb
@@ -11,7 +13,7 @@ from pytorch_lightning.loggers import WandbLogger
 from pytorch_lightning.callbacks import ModelCheckpoint
 from torch.nn import MSELoss
 from SDOAIA_dataloader import AIA_GOESDataModule
-from flaring.MEGS_AI_baseline.models.linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel
+from models.linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel
 from callback import ImagePredictionLogger_SXR
 
 # Parser
@@ -115,7 +117,7 @@ for parameter_set in combined_parameters:
         default_root_dir=checkpoint_dir,
         accelerator="gpu" if torch.cuda.is_available() else "cpu",
         devices=1,
-        max_epochs=run_config.get('epochs', 10),
+        max_epochs=run_config['epochs'],
         callbacks=[sxr_plot_callback, checkpoint_callback],
         logger=wandb_logger,
         log_every_n_steps=10
@@ -131,9 +133,15 @@ for parameter_set in combined_parameters:
     full_checkpoint_path = os.path.join(checkpoint_dir, f"{wb_name}_{n}.ckpt")
     torch.save(save_dictionary, full_checkpoint_path)
 
-    # Test
-    trainer.test(model, dataloaders=data_loader.test_dataloader())
-
+    # Save final PyTorch checkpoint with model and state_dict
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    final_checkpoint_path = os.path.join(checkpoint_dir, f"{wb_name}-final-{timestamp}.pth")
+    torch.save({
+        'model': model,
+        'state_dict': model.state_dict()
+    }, final_checkpoint_path)
+    print(f"Saved final PyTorch checkpoint: {final_checkpoint_path}")
+    n += 1
     # Finalize
     wandb.finish()
     n += 1