changed structure

flaring/MEGS_AI_baseline/SDOAIA_dataloader.py

@@ -57,19 +57,19 @@ class AIA_GOESDataset(torch.utils.data.Dataset):
                                      self.target_size[0]/aia_img.shape[1],
                                      self.target_size[1]/aia_img.shape[2]))
 
-       #Apply cut and normalize:
-        cuts_dict = {
-            0: np.float32(16.560747),
-            1: np.float32(75.84181),
-            2: np.float32(1536.1443),
-            3: np.float32(2288.1),
-            4: np.float32(1163.9178),
-            5: np.float32(401.82352)
-        }
-
-        for channel in range(6):
-            aia_img[channel] = np.clip(aia_img[channel], 0, cuts_dict[channel])
-            aia_img[channel] = aia_img[channel] / cuts_dict[channel]  # Normalize each channel to [0, 1]
+       # #Apply cut and normalize:
+       #  cuts_dict = {
+       #      0: np.float32(16.560747),
+       #      1: np.float32(75.84181),
+       #      2: np.float32(1536.1443),
+       #      3: np.float32(2288.1),
+       #      4: np.float32(1163.9178),
+       #      5: np.float32(401.82352)
+       #  }
+       #
+       #  for channel in range(6):
+       #      aia_img[channel] = np.clip(aia_img[channel], 0, cuts_dict[channel])
+       #      aia_img[channel] = aia_img[channel] / cuts_dict[channel]  # Normalize each channel to [0, 1]
 
         # Convert to torch for transforms
         aia_img = torch.tensor(aia_img, dtype=torch.float32) # (6, H, W)

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,42 @@ 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 unnormalize_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)
 
-            aia_imgs = aia_imgs.to(next(model.parameters()).device)
-            log_outputs = model(aia_imgs)  # Get raw log-space outputs
+            # Get model predictions
+            log_outputs = model(aia_imgs)
+
+            # 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 data')
+    parser.add_argument('--sxr_norm', required=True, help='Path to SXR normalization parameters (mean, std)')
     parser.add_argument('--output', default='log_predictions.txt',
                         help='Output file for log-space predictions')
     parser.add_argument('--batch-size', type=int, default=8,
@@ -32,23 +46,42 @@ def main():
 
     args = parser.parse_args()
 
+    sxr_norm = np.load(args.sxr_norm)
+
     # 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,
+<<<<<<< HEAD
         sxr_dir='',
         sxr_norm=None,
         transform=None
+=======
+        sxr_dir=args.sxr_dir,  # No SXR files needed
+        transform=None  # No input transforms
+>>>>>>> 22f4a17192a3a77fa4d4fe1ae3a2aa8c0bbdb539
     )
 
     # Save log-space predictions
     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
+            pred = unnormalize_sxr(log_pred, sxr_norm)
+            print(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/{linear_and_hybrid.py → 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,19 +1,19 @@
 
 import argparse
 import os
+from datetime import datetime
+
 import yaml
 import itertools
 import wandb
 import torch
 import numpy as np
-from pathlib import Path
-import torchvision.transforms as transforms
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import WandbLogger
-from pytorch_lightning.callbacks import ModelCheckpoint, Callback
-from torch.nn import HuberLoss, MSELoss
+from pytorch_lightning.callbacks import ModelCheckpoint
+from torch.nn import MSELoss
 from SDOAIA_dataloader import AIA_GOESDataModule
-from linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel
+from models.linear_and_hybrid import LinearIrradianceModel, HybridIrradianceModel
 from callback import ImagePredictionLogger_SXR
 
 # Parser
@@ -117,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
@@ -127,15 +127,16 @@ for parameter_set in combined_parameters:
     trainer.fit(model, data_loader)
 
     # Save checkpoint
-    save_dictionary = run_config
-    save_dictionary['model'] = model
-    save_dictionary['instrument'] = instrument
-    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())
+    trainer.fit(model, data_loader)
 
+    # 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
+    wandb.finish()

flaring/normalization_and_aligning_data.py

@@ -6,11 +6,13 @@ from astropy.io import fits
 import warnings
 import pandas as pd
 from astropy.visualization import ImageNormalize, AsinhStretch
+from multiprocessing import Pool, cpu_count
+from functools import partial
+import time
+from tqdm import tqdm
 
 warnings.filterwarnings('ignore')
 
-import pandas as pd
-
 # Directory paths for each wavelength folder.
 wavelength_dirs = {
     "94": "/mnt/data2/AIA_processed_data/94",
@@ -22,50 +24,8 @@ wavelength_dirs = {
 }
 
 # Regular expression to extract timestamp from file names.
-# Adjust this pattern to match your file naming scheme.
 timestamp_pattern = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}")
 
-# Collect timestamps found in each wavelength directory.
-timestamps_found = defaultdict(set)
-
-for wavelength, dir_path in wavelength_dirs.items():
-    try:
-        for filename in os.listdir(dir_path):
-            match = timestamp_pattern.search(filename)
-            if match:
-                ts = match.group(0)
-                timestamps_found[ts].add(wavelength)
-    except Exception as e:
-        print(f"Could not read directory {dir_path}: {e}")
-
-# Identify timestamps that exist in all wavelength folders.
-all_wavelengths = set(wavelength_dirs.keys())
-common_timestamps = [ts for ts, waves in timestamps_found.items() if waves == all_wavelengths]
-
-# Identify which timestamps are missing files for some wavelengths.
-missing_files = {
-    ts: list(all_wavelengths - waves)
-    for ts, waves in timestamps_found.items() if waves != all_wavelengths
-}
-
-print("Timestamps present in all wavelength folders:")
-for ts in sorted(common_timestamps):
-    print(ts)
-
-print("\nTimestamps with missing wavelength files:")
-for ts, missing in missing_files.items():
-    print(f"{ts}: missing {', '.join(sorted(missing))}")
-
-
-goes = pd.read_csv("/mnt/data/goes_combined/combined_g18_avg1m_20230701_20230815.csv")
-# Convert 'time' column to datetime
-goes['time'] = pd.to_datetime(goes['time'], format='%Y-%m-%d %H:%M:%S')
-
-
-# Initialize the array to store all wavelength data
-data_shape = (6, 512, 512)
-
-
 # Map wavelengths to array indices
 wavelength_to_idx = {
     '94': 0,
@@ -76,52 +36,171 @@ wavelength_to_idx = {
     '304': 5
 }
 
-sdo_norms = {0: ImageNormalize(vmin=0, vmax= np.float32(16.560747), stretch=AsinhStretch(0.005), clip=True),
-             1: ImageNormalize(vmin=0, vmax= np.float32(75.84181), stretch=AsinhStretch(0.005), clip=True),
-             2: ImageNormalize(vmin=0, vmax= np.float32(1536.1443), stretch=AsinhStretch(0.005), clip=True),
-             3: ImageNormalize(vmin=0, vmax= np.float32(2288.1), stretch=AsinhStretch(0.005), clip=True),
-             4: ImageNormalize(vmin=0, vmax=np.float32(1163.9178), stretch=AsinhStretch(0.005), clip=True),
-             5: ImageNormalize(vmin=0, vmax=np.float32(401.82352), stretch=AsinhStretch(0.001), clip=True),
-             }
-
-
-
-# Load data for each timestamp and wavelength
-for time_idx, timestamp in enumerate(common_timestamps):
-    sxr = goes[goes['time'] == pd.to_datetime(timestamp)]
-    sxr_a = sxr['xrsa_flux'].values[0] if not sxr.empty else None
-    sxr_b = sxr['xrsb_flux'].values[0] if not sxr.empty else None
-    if sxr_a is None or sxr_b is None:
-        print(f"Missing SXR data for timestamp {timestamp}, skipping...")
-        continue
-    wavelength_data = np.zeros(data_shape, dtype=np.float32)
-    sxr_a_data = np.zeros(1, dtype=np.float32)
-    sxr_b_data = np.zeros(1, dtype=np.float32)
-    sxr_a_data[0] = sxr_a if sxr_a is not None else np.nan
-    sxr_b_data[0] = sxr_b if sxr_b is not None else np.nan
-    print(f"Processing timestamp: {timestamp} (Index: {time_idx})")
-    for wavelength, wave_idx in wavelength_to_idx.items():
-        filepath = os.path.join(wavelength_dirs[wavelength], f"{timestamp}.fits")
-        with fits.open(filepath) as hdul:
-            raw_data = hdul[0].data
-
-            # Apply the appropriate normalization for this wavelength
-            if wave_idx in sdo_norms:
-                # Get the normalizer for this wavelength index
-                normalizer = sdo_norms[wave_idx]
-
-                # Apply normalization and convert to [-1, 1] range
-                normalized_data = normalizer(raw_data)
-                wavelength_data[wave_idx] = normalized_data * 2 - 1
-            else:
-                # Fallback if no normalizer exists for this wavelength
-                print(f"Warning: No normalizer found for wavelength index {wave_idx}")
-                wavelength_data[wave_idx] = raw_data
-
-    # Store the wavelength data for this timestamp
-    np.save(f"/mnt/data2/ML-Ready/AIA-Data/{timestamp}.npy", wavelength_data)
-    # Store the SXR data
-    np.save(f"/mnt/data2/ML-Ready/GOES-18-SXR-A/{timestamp}.npy", sxr_a_data)
-    np.save(f"/mnt/data2/ML-Ready/GOES-18-SXR-B/{timestamp}.npy", sxr_b_data)
-    print(f"Saved data for timestamp {timestamp} to disk.")
-    print(f"Percent: {time_idx + 1} / {len(common_timestamps)}")
+# Initialize the array to store all wavelength data
+data_shape = (6, 512, 512)
+
+sdo_norms = {
+    0: ImageNormalize(vmin=0, vmax=np.float32(16.560747), stretch=AsinhStretch(0.005), clip=True),
+    1: ImageNormalize(vmin=0, vmax=np.float32(75.84181), stretch=AsinhStretch(0.005), clip=True),
+    2: ImageNormalize(vmin=0, vmax=np.float32(1536.1443), stretch=AsinhStretch(0.005), clip=True),
+    3: ImageNormalize(vmin=0, vmax=np.float32(2288.1), stretch=AsinhStretch(0.005), clip=True),
+    4: ImageNormalize(vmin=0, vmax=np.float32(1163.9178), stretch=AsinhStretch(0.005), clip=True),
+    5: ImageNormalize(vmin=0, vmax=np.float32(401.82352), stretch=AsinhStretch(0.001), clip=True),
+}
+
+
+def process_timestamp(args):
+    """
+    Process a single timestamp: load wavelength data, apply normalization,
+    and save to disk along with SXR data.
+    """
+    timestamp, goes_data = args
+    try:
+        # Get SXR data for this timestamp
+        sxr = goes_data[goes_data['time'] == pd.to_datetime(timestamp)]
+        sxr_a = sxr['xrsa_flux'].values[0] if not sxr.empty else None
+        sxr_b = sxr['xrsb_flux'].values[0] if not sxr.empty else None
+
+        if sxr_a is None or sxr_b is None:
+            return (timestamp, False, f"Missing SXR data for timestamp {timestamp}")
+
+        # Initialize arrays
+        wavelength_data = np.zeros(data_shape, dtype=np.float32)
+        sxr_a_data = np.zeros(1, dtype=np.float32)
+        sxr_b_data = np.zeros(1, dtype=np.float32)
+        sxr_a_data[0] = sxr_a
+        sxr_b_data[0] = sxr_b
+
+        # Process each wavelength
+        for wavelength, wave_idx in wavelength_to_idx.items():
+            filepath = os.path.join(wavelength_dirs[wavelength], f"{timestamp}.fits")
+
+            with fits.open(filepath) as hdul:
+                raw_data = hdul[0].data
+
+                # Apply the appropriate normalization for this wavelength
+                if wave_idx in sdo_norms:
+                    normalizer = sdo_norms[wave_idx]
+                    normalized_data = normalizer(raw_data)
+                    wavelength_data[wave_idx] = normalized_data * 2 - 1
+                else:
+                    wavelength_data[wave_idx] = raw_data
+
+        # Save data to disk
+        np.save(f"/mnt/data2/ML-Ready/AIA-Data/{timestamp}.npy", wavelength_data)
+        np.save(f"/mnt/data2/ML-Ready/GOES-18-SXR-A/{timestamp}.npy", sxr_a_data)
+        np.save(f"/mnt/data2/ML-Ready/GOES-18-SXR-B/{timestamp}.npy", sxr_b_data)
+
+        return (timestamp, True, "Success")
+
+    except Exception as e:
+        return (timestamp, False, f"Error processing timestamp {timestamp}: {e}")
+
+
+def update_progress(result):
+    """Callback function to update progress bar"""
+    global pbar, successful_count, failed_count
+    timestamp, success, message = result
+
+    if success:
+        successful_count += 1
+        pbar.set_postfix(success=successful_count, failed=failed_count)
+    else:
+        failed_count += 1
+        pbar.set_postfix(success=successful_count, failed=failed_count)
+        tqdm.write(f"Failed: {message}")
+
+    pbar.update(1)
+
+
+def main():
+    global pbar, successful_count, failed_count
+
+    # Collect timestamps found in each wavelength directory.
+    timestamps_found = defaultdict(set)
+
+    print("Scanning directories for timestamps...")
+    for wavelength, dir_path in tqdm(wavelength_dirs.items(), desc="Scanning directories"):
+        try:
+            for filename in os.listdir(dir_path):
+                match = timestamp_pattern.search(filename)
+                if match:
+                    ts = match.group(0)
+                    timestamps_found[ts].add(wavelength)
+        except Exception as e:
+            print(f"Could not read directory {dir_path}: {e}")
+
+    # Identify timestamps that exist in all wavelength folders.
+    all_wavelengths = set(wavelength_dirs.keys())
+    common_timestamps = [ts for ts, waves in timestamps_found.items() if waves == all_wavelengths]
+
+    # Identify which timestamps are missing files for some wavelengths.
+    missing_files = {
+        ts: list(all_wavelengths - waves)
+        for ts, waves in timestamps_found.items() if waves != all_wavelengths
+    }
+
+    print(f"\nFound {len(common_timestamps)} timestamps present in all wavelength folders")
+    print(f"Found {len(missing_files)} timestamps with missing wavelength files")
+
+    # Load GOES data
+    print("Loading GOES data...")
+    goes = pd.read_csv("/mnt/data/goes_combined/combined_g18_avg1m_20230701_20230815.csv")
+    goes['time'] = pd.to_datetime(goes['time'], format='%Y-%m-%d %H:%M:%S')
+
+    # Create output directories if they don't exist
+    os.makedirs("/mnt/data2/ML-Ready/AIA-Data", exist_ok=True)
+    os.makedirs("/mnt/data2/ML-Ready/GOES-18-SXR-A", exist_ok=True)
+    os.makedirs("/mnt/data2/ML-Ready/GOES-18-SXR-B", exist_ok=True)
+
+    # Use all available CPU cores
+    num_processes = cpu_count()
+    print(f"Using {num_processes} CPU cores for processing")
+    print(f"Processing {len(common_timestamps)} timestamps...")
+
+    # Initialize global counters for progress tracking
+    successful_count = 0
+    failed_count = 0
+
+    # Create arguments for multiprocessing (timestamp, goes_data pairs)
+    args_list = [(timestamp, goes) for timestamp in common_timestamps]
+
+    # Start timing
+    start_time = time.time()
+
+    # Create progress bar
+    pbar = tqdm(total=len(common_timestamps), desc="Processing timestamps",
+                unit="timestamp", dynamic_ncols=True)
+
+    # Process timestamps in parallel with progress tracking
+    with Pool(processes=num_processes) as pool:
+        # Use map with callback for real-time progress updates
+        results = []
+        for args in args_list:
+            result = pool.apply_async(process_timestamp, (args,), callback=update_progress)
+            results.append(result)
+
+        # Wait for all processes to complete
+        for result in results:
+            result.wait()
+
+    # Close progress bar
+    pbar.close()
+
+    # Calculate statistics
+    end_time = time.time()
+    total_time = end_time - start_time
+
+    print(f"\nProcessing complete!")
+    print(f"Total time: {total_time:.2f} seconds")
+    print(f"Average time per timestamp: {total_time / len(common_timestamps):.2f} seconds")
+    print(f"Successfully processed: {successful_count}/{len(common_timestamps)} timestamps")
+    print(f"Failed processes: {failed_count}")
+    print(f"Processing rate: {len(common_timestamps) / total_time:.2f} timestamps/second")
+
+    if failed_count > 0:
+        print(f"\n{failed_count} timestamps failed processing (see messages above)")
+
+
+if __name__ == "__main__":
+    main()

flaring/split_data.py

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 import os
 import pandas as pd
 import shutil
@@ -75,9 +76,94 @@ for base_dir in [flares_event_dir, non_flares_event_dir]:
         else:
             print(f"Skipping file {file} in {base_dir}: Outside date range")
             continue
+=======
+#
+
+
+#
+# data_dir = "/mnt/data/ML-Ready/AIA-Data"
+# flares_event_dir = "/mnt/data/ML-Ready/flares_event_dir"
+# non_flares_event_dir = "/mnt/data/ML-Ready/non_flares_event_dir"
+# flare_events_csv = "/mnt/data/flare_list/flare_events_2023-07-01_2023-08-15.csv"
+#
+# train_range = (datetime(2023, 7, 1), datetime(2023, 7, 25))
+# val_range = (datetime(2023, 7, 27), datetime(2023, 7, 30))
+# test_range = (datetime(2023, 8, 1), datetime(2023, 8, 15))
+#
+# os.makedirs(flares_event_dir, exist_ok=True)
+# os.makedirs(non_flares_event_dir, exist_ok=True)
+#
+# os.makedirs(os.path.join(flares_event_dir, "train"), exist_ok=True)
+# os.makedirs(os.path.join(flares_event_dir, "val"), exist_ok=True)
+# os.makedirs(os.path.join(flares_event_dir, "test"), exist_ok=True)
+#
+# os.makedirs(os.path.join(non_flares_event_dir, "train"), exist_ok=True)
+# os.makedirs(os.path.join(non_flares_event_dir, "val"), exist_ok=True)
+# os.makedirs(os.path.join(non_flares_event_dir, "test"), exist_ok=True)
+#
+#
+# flare_event = pd.read_csv(flare_events_csv)
+# print(f"Found {len(flare_event)} flare events")
+# 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))
+#
+# data_list = os.listdir(data_dir)
+# print(f"Found {len(data_list)} files in {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
+#
+#     # Check if the file's time falls within any flare event
+#     is_flaring = any(start <= aia_time <= end for start, end in flaring_eve_list)
+#     if is_flaring:
+#         src = os.path.join(data_dir, file)
+#         dst = os.path.join(flares_event_dir, file)
+#
+#         if train_range[0] <= aia_time <= train_range[1]:
+#             dst = os.path.join(flares_event_dir, "train")
+#             shutil.copy(src, dst)
+#         elif val_range[0] <= aia_time <= val_range[1]:
+#             dst = os.path.join(flares_event_dir, "val")
+#             shutil.copy(src, dst)
+#         elif test_range[0] <= aia_time <= test_range[1]:
+#             dst = os.path.join(flares_event_dir, "test")
+#             shutil.copy(src, dst)
+#         else:
+#             print(f"Skipping {file}: Time {aia_time} not in any defined range")
+#             continue
+#         print(f"Copied {file} to {dst}")
+#     else:
+#         print("Skipping non-flaring event file:", file)
+    # else:
+    #     src = os.path.join(data_dir, file)
+    #     dst = os.path.join(non_flares_event_dir, file)
+    #     print(aia_time)
+    #     print(train_range[0], train_range[1])
+    #     if train_range[0] <= aia_time <= train_range[1]:
+    #         split_dir = "train"
+    #     elif val_range[0] <= aia_time <= val_range[1]:
+    #         split_dir = "val"
+    #     elif test_range[0] <= aia_time <= test_range[1]:
+    #         split_dir = "test"
+    #     dst = os.path.join(flares_event_dir, split_dir)
+    #    # shutil.copy(src, dst)
+    #     print(f"Copied {file} to {dst}")
+
+
+>>>>>>> 22f4a17192a3a77fa4d4fe1ae3a2aa8c0bbdb539
 
         # Move file to appropriate split directory
         src = os.path.join(base_dir, file)
         dst = os.path.join(base_dir, split_dir, file)
         shutil.move(src, dst)
-        print(f"Moved {file} to {base_dir}/{split_dir}")
+<<<<<<< HEAD
+        print(f"Moved {file} to {base_dir}/{split_dir}")
+=======
+        print(f"Moved {file} to {base_dir}/{split_dir}")
+>>>>>>> 22f4a17192a3a77fa4d4fe1ae3a2aa8c0bbdb539