Add ablation study and spatial performance scripts; update pipeline configuration and dataset handling

.gitignore

@@ -155,4 +155,11 @@ wandb/
 
 .claude/
 Untracked/
-foxes.yml
+foxes.yml
+.git-backup
+/forecasting/training/convert_flux_to_npy.py
+/forecasting/inference/extract_patch_flux.py
+/.merged_train_config.yaml
+/.merged_inference_config.yaml
+/.merged_evaluate_config.yaml
+.npz

analysis/spatial_performance.py

@@ -0,0 +1,273 @@
+"""
+Flux-Weighted Error Heatmap on Solar Disk
+==========================================
+For each matched flux map, accumulates:
+    mae_sum[i,j]  += flux[i,j] * |log10 error|
+    bias_sum[i,j] += flux[i,j] *  log10 error
+    weight[i,j]   += flux[i,j]
+
+Then normalizes to get flux-weighted mean error per patch.
+
+Usage
+-----
+    python analysis/spatial_performance.py
+
+Outputs
+-------
+    analysis/flux_weighted_errors_t0.npz     — accumulation cache
+    analysis/performance_heatmap_all.png
+"""
+
+import os
+import sys
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tqdm import tqdm
+from pathlib import Path
+from cmap import Colormap
+
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+from forecasting.inference.evaluation import setup_barlow_font
+
+# ---------------------------------------------------------------------------
+# Paths — edit here
+# ---------------------------------------------------------------------------
+FLUX_DIR        = "/Volumes/T9/FOXES_Data/flux/"
+PREDICTIONS_CSV = "/Volumes/T9/FOXES_Misc/batch_results/vit/vit_predictions_test.csv"
+OUT_DIR         = Path(__file__).parent
+GRID_SIZE       = 64    # 512px / 8px patch size
+BIN_SIZE        = 1     # downsample factor (1 = full 64×64 resolution)
+CROP_FACTOR     = 1.1   # AIA images cropped at 1.1 solar radii
+SOLAR_RADIUS_PATCHES = (GRID_SIZE / 2) / CROP_FACTOR   # ≈ 29.1 patches
+
+# Only patches above this percentile (per flux map) contribute.
+# 0 = include all non-zero patches.
+FLUX_THRESHOLD_PERCENTILE = 0
+
+# Percentile cap for colorbar scaling (applied to non-NaN values).
+# e.g. 99 clips the top 1% of values so detail in the bulk is visible.
+VMAX_PERCENTILE = 99.9
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def normalize_ts(series: pd.Series) -> pd.Series:
+    return pd.to_datetime(
+        series.astype(str).str.replace("_", ":", regex=False), utc=False,
+    ).dt.floor("s")
+
+
+def _ts_key(fpath: str) -> str:
+    raw = os.path.basename(fpath).replace('.npy', '').replace('_', ':')
+    return pd.Timestamp(raw).floor('s').isoformat()
+
+
+def load_predictions(predictions_csv: str) -> pd.DataFrame:
+    df = pd.read_csv(predictions_csv)
+    df["timestamp"] = normalize_ts(df["timestamp"])
+    df["log_pred"]      = np.log10(df["predictions"])
+    df["log_gt"]        = np.log10(df["groundtruth"])
+    df["log_error"]     = df["log_pred"] - df["log_gt"]
+    df["log_abs_error"] = df["log_error"].abs()
+    print(f"Loaded {len(df)} predictions")
+    return df
+
+
+# ---------------------------------------------------------------------------
+# Heatmap accumulation
+# ---------------------------------------------------------------------------
+
+# NOTE: module-level for ProcessPoolExecutor (spawn on macOS)
+def _accumulate_flux_map(args):
+    fpath, log_abs_error, log_error, threshold_pct = args
+    fmap = np.load(fpath).astype(np.float64)
+    active = fmap[fmap > 0]
+    if active.size == 0:
+        return None
+    if threshold_pct > 0:
+        thresh = np.percentile(active, threshold_pct)
+        fmap = np.where(fmap >= thresh, fmap, 0.0)
+    else:
+        fmap = np.where(fmap > 0, fmap, 0.0)
+    total = fmap.sum()
+    if total == 0:
+        return None
+    fmap = fmap / total  # normalise: each timestamp contributes equal total weight
+    return fmap * log_abs_error, fmap * log_error, fmap
+
+
+def compute_flux_weighted_errors(flux_dir: str, df: pd.DataFrame, cache_path: Path,
+                                  threshold_pct: int = FLUX_THRESHOLD_PERCENTILE) -> dict:
+    cache_path = cache_path.with_stem(f"{cache_path.stem}_t{threshold_pct}")
+
+    if cache_path.exists():
+        print(f"Loading cached flux-weighted error maps from {cache_path}")
+        data = np.load(cache_path)
+        result = {}
+        n = float(data['count'])
+        w = data['weight']
+        mae  = np.where(w > 0, data['mae_sum']  / n, np.nan) if n > 0 else np.full_like(w, np.nan)
+        bias = np.where(w > 0, data['bias_sum'] / n, np.nan) if n > 0 else np.full_like(w, np.nan)
+        return mae, bias, w
+
+    lookup = {}
+    for _, row in df.iterrows():
+        key = pd.Timestamp(row['timestamp']).floor('s').isoformat()
+        lookup[key] = (float(row['log_abs_error']), float(row['log_error']))
+
+    shape = (GRID_SIZE, GRID_SIZE)
+    mae_sum  = np.zeros(shape)
+    bias_sum = np.zeros(shape)
+    weight   = np.zeros(shape)
+    count    = 0
+
+    files = sorted([os.path.join(flux_dir, f)
+                    for f in os.listdir(flux_dir) if f.endswith('.npy')])
+    args_list = []
+    for fpath in files:
+        try:
+            ts_key = _ts_key(fpath)
+        except Exception:
+            continue
+        if ts_key not in lookup:
+            continue
+        abs_err, err = lookup[ts_key]
+        args_list.append((fpath, abs_err, err, threshold_pct))
+
+    print(f"Matched {len(args_list)} / {len(files)} flux maps")
+
+    with ProcessPoolExecutor(max_workers=os.cpu_count()) as executor:
+        futures = {executor.submit(_accumulate_flux_map, a): i
+                   for i, a in enumerate(args_list)}
+        for future in tqdm(as_completed(futures), total=len(args_list),
+                           desc="Accumulating flux-weighted errors"):
+            result = future.result()
+            if result is None:
+                continue
+            mae_c, bias_c, flux_c = result
+            mae_sum  += mae_c
+            bias_sum += bias_c
+            weight   += flux_c
+            count    += 1
+
+    np.savez(cache_path, mae_sum=mae_sum, bias_sum=bias_sum,
+             weight=weight, count=np.array(count))
+    print(f"Saved → {cache_path}")
+
+    mae  = np.where(weight > 0, mae_sum  / count, np.nan) if count > 0 else np.full(shape, np.nan)
+    bias = np.where(weight > 0, bias_sum / count, np.nan) if count > 0 else np.full(shape, np.nan)
+    return mae, bias, weight
+
+
+# ---------------------------------------------------------------------------
+# Plot
+# ---------------------------------------------------------------------------
+
+def _bin_grid(grid: np.ndarray, bin_size: int) -> np.ndarray:
+    if bin_size == 1:
+        return grid
+    h, w = grid.shape
+    bh, bw = h // bin_size, w // bin_size
+    cropped = grid[:bh * bin_size, :bw * bin_size]
+    return np.nanmean(cropped.reshape(bh, bin_size, bw, bin_size), axis=(1, 3))
+
+
+def plot_flux_weighted_heatmap(mae_grid: np.ndarray, bias_grid: np.ndarray,
+                                weight_grid: np.ndarray, out_path: Path,
+                                subtitle: str = "", bin_size: int = BIN_SIZE,
+                                vmax_pct: int = VMAX_PERCENTILE):
+    setup_barlow_font()
+    text_color = "#111111"
+    theta      = np.linspace(0, 2 * np.pi, 300)
+
+    mae_b    = _bin_grid(mae_grid, bin_size)
+    bias_b   = _bin_grid(bias_grid, bin_size)
+    weight_b = _bin_grid(weight_grid, bin_size)
+    n_bins   = mae_b.shape[0]
+    cy, cx   = n_bins / 2, n_bins / 2
+
+    r_limb       = SOLAR_RADIUS_PATCHES / bin_size
+
+    mae_vmax  = np.nanpercentile(mae_b, vmax_pct)
+    mae_norm  = plt.Normalize(vmin=0, vmax=mae_vmax)
+
+    bias_cap  = np.nanpercentile(np.abs(bias_b), vmax_pct)
+    bias_norm = plt.Normalize(vmin=-bias_cap, vmax=bias_cap)
+
+    panels = [
+        (mae_b,   r"Flux-Weighted MAE (log$_{10}$ Space)",  Colormap('cmocean:thermal').to_mpl(),  mae_norm),
+        (bias_b,  r"Flux-Weighted MBE (log$_{10}$ Space)",  Colormap('cmasher:fusion_r').to_mpl(), bias_norm),
+        # (np.log10(np.where(weight_b > 0, weight_b, np.nan)),
+        #           r"log$_{10}$ Accumulated Flux",                 "viridis", None),
+    ]
+
+    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
+    fig.patch.set_facecolor("white")
+
+    for ax, (grid, title, cmap, norm) in zip(axes, panels):
+        im = ax.imshow(grid, origin="lower", cmap=cmap, norm=norm,
+                       interpolation="bilinear", extent=[0, n_bins, 0, n_bins])
+        cbar = fig.colorbar(im, ax=ax, shrink=0.82)
+        cbar.ax.tick_params(labelsize=9, colors=text_color)
+        def _fmt(x, _):
+            m, e = f"{x:.2e}".split("e")
+            return f"{m}e{int(e)}"
+        cbar.ax.yaxis.set_major_formatter(plt.matplotlib.ticker.FuncFormatter(_fmt))
+        for lbl in cbar.ax.get_yticklabels():
+            lbl.set_fontfamily("Barlow")
+            lbl.set_fontsize(9)
+            lbl.set_color(text_color)
+        #cbar.set_label(title, fontsize=9, color=text_color, fontfamily="Barlow")
+
+        ax.plot(cx + r_limb       * np.cos(theta), cy + r_limb       * np.sin(theta),
+                color="#4488FF", linestyle="--", linewidth=1.2, alpha=0.8,
+                label=f"Solar Limb")
+
+        tick_bins   = np.linspace(0, n_bins, 7)
+        tick_labels = [f"{int((t - n_bins / 2) * bin_size)}" for t in tick_bins]
+        ax.set_xticks(tick_bins); ax.set_xticklabels(tick_labels)
+        ax.set_yticks(tick_bins); ax.set_yticklabels(tick_labels)
+
+        ax.set_title(title, fontsize=10, color=text_color, fontfamily="Barlow")
+        ax.set_xlabel("Solar X (ViT Patches From Center)", fontsize=9,
+                      color=text_color, fontfamily="Barlow")
+        ax.set_ylabel("Solar Y (ViT Patches From Center)", fontsize=9,
+                      color=text_color, fontfamily="Barlow")
+        ax.tick_params(labelsize=8, colors=text_color)
+        ax.legend(fontsize=7, facecolor="white", edgecolor="grey", loc="upper right",)
+        for spine in ax.spines.values():
+            spine.set_color(text_color)
+
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=400, bbox_inches="tight", facecolor="white")
+    plt.show()
+    print(f"Saved → {out_path}")
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--flux_dir",        default=FLUX_DIR)
+    parser.add_argument("--predictions_csv", default=PREDICTIONS_CSV)
+    parser.add_argument("--out_dir",         default=str(OUT_DIR))
+    args = parser.parse_args()
+
+    out = Path(args.out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    df = load_predictions(args.predictions_csv)
+
+    mae, bias, weight = compute_flux_weighted_errors(
+        args.flux_dir, df, out / "flux_weighted_errors.npz"
+    )
+    plot_flux_weighted_heatmap(mae, bias, weight,
+                               out / "performance_heatmap_all.png",
+                               subtitle="All flares")

forecasting/data_loaders/SDOAIA_dataloader.py

@@ -63,7 +63,7 @@ class AIA_GOESDataset(torch.utils.data.Dataset):
         Strategy for balancing dataset classes.
     """
 
-    def __init__(self, aia_dir, sxr_dir, wavelengths=[94, 131, 171, 193, 211, 304], sxr_transform=None,
+    def __init__(self, aia_dir, sxr_dir, wavelengths=[94, 131, 171, 193, 211, 304, 335], sxr_transform=None,
                  target_size=(512, 512), cadence=1, reference_time=None, only_prediction=False, oversample=False,
                  flare_threshold=1e-5, balance_strategy='upsample_minority'):
         self.aia_dir = Path(aia_dir).resolve()
@@ -311,6 +311,130 @@ class AIA_GOESDataset(torch.utils.data.Dataset):
         return timestamp
 
 
+class NoisyAIA_GOESDataset(AIA_GOESDataset):
+    """
+    Ablation dataset that applies Gaussian noise to specific AIA wavelength channels.
+
+    Inherits all behavior from AIA_GOESDataset and injects additive Gaussian noise
+    on the specified wavelengths after loading, enabling channel-wise ablation studies.
+
+    Parameters
+    ----------
+    noisy_wavelengths : list of int
+        Subset of wavelengths to corrupt with noise (e.g. [171, 193]).
+        Must be a subset of the `wavelengths` argument.
+    noise_std : float or dict
+        Standard deviation of the Gaussian noise. If a float, the same std is
+        applied to all noisy channels. If a dict mapping wavelength -> std, each
+        channel gets its own scale.
+    *args, **kwargs
+        Forwarded to AIA_GOESDataset.
+    """
+
+    def __init__(self, *args, noisy_wavelengths, noise_std=1.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.noisy_wavelengths = noisy_wavelengths
+        # Build a per-channel std lookup
+        if isinstance(noise_std, dict):
+            self.noise_std = noise_std
+        else:
+            self.noise_std = {w: noise_std for w in noisy_wavelengths}
+
+        invalid = set(noisy_wavelengths) - set(self.wavelengths)
+        if invalid:
+            raise ValueError(f"noisy_wavelengths {invalid} not in dataset wavelengths {self.wavelengths}")
+
+    def __getitem__(self, idx):
+        aia_img, sxr_val = super().__getitem__(idx)  # (H, W, C)
+
+        for wav in self.noisy_wavelengths:
+            c = self.wavelengths.index(wav)
+            std = self.noise_std[wav]
+            noise = torch.randn_like(aia_img[..., c]) * std
+            aia_img[..., c] = aia_img[..., c] + noise
+
+        return aia_img, sxr_val
+
+
+class NoisyAIA_GOESDataModule(LightningDataModule):
+    """
+    DataModule for the Gaussian-noise ablation study.
+
+    Wraps NoisyAIA_GOESDataset for test/val splits while keeping training
+    clean (or optionally noisy too). Noise is applied only at evaluation time
+    by default so the pre-trained model is evaluated without retraining.
+
+    Parameters
+    ----------
+    noisy_wavelengths : list of int
+        Wavelengths to corrupt during evaluation.
+    noise_std : float or dict
+        Noise standard deviation(s); see NoisyAIA_GOESDataset.
+    apply_noise_to_train : bool, optional
+        If True, also corrupt the training split (default: False).
+    All other parameters match AIA_GOESDataModule.
+    """
+
+    def __init__(self, aia_train_dir, aia_val_dir, aia_test_dir,
+                 sxr_train_dir, sxr_val_dir, sxr_test_dir,
+                 sxr_norm, noisy_wavelengths, noise_std=1.0,
+                 apply_noise_to_train=False,
+                 batch_size=64, num_workers=4,
+                 wavelengths=[94, 131, 171, 193, 211, 304, 335],
+                 cadence=1, reference_time=None):
+        super().__init__()
+        self.aia_train_dir = aia_train_dir
+        self.aia_val_dir = aia_val_dir
+        self.aia_test_dir = aia_test_dir
+        self.sxr_train_dir = sxr_train_dir
+        self.sxr_val_dir = sxr_val_dir
+        self.sxr_test_dir = sxr_test_dir
+        self.sxr_norm = sxr_norm
+        self.noisy_wavelengths = noisy_wavelengths
+        self.noise_std = noise_std
+        self.apply_noise_to_train = apply_noise_to_train
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.wavelengths = wavelengths
+        self.cadence = cadence
+        self.reference_time = reference_time
+
+    def _make_dataset(self, aia_dir, sxr_dir, noisy):
+        cls = NoisyAIA_GOESDataset if noisy else AIA_GOESDataset
+        kwargs = dict(
+            aia_dir=aia_dir,
+            sxr_dir=sxr_dir,
+            sxr_transform=SXRLogNormTransform(self.sxr_norm[0], self.sxr_norm[1]),
+            target_size=(512, 512),
+            wavelengths=self.wavelengths,
+            cadence=self.cadence,
+            reference_time=self.reference_time,
+        )
+        if noisy:
+            kwargs["noisy_wavelengths"] = self.noisy_wavelengths
+            kwargs["noise_std"] = self.noise_std
+        return cls(**kwargs)
+
+    def setup(self, stage=None):
+        self.train_ds = self._make_dataset(
+            self.aia_train_dir, self.sxr_train_dir, noisy=self.apply_noise_to_train
+        )
+        self.val_ds = self._make_dataset(self.aia_val_dir, self.sxr_val_dir, noisy=True)
+        self.test_ds = self._make_dataset(self.aia_test_dir, self.sxr_test_dir, noisy=True)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_ds, batch_size=self.batch_size,
+                          shuffle=True, num_workers=self.num_workers, prefetch_factor=4)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_ds, batch_size=self.batch_size,
+                          shuffle=False, num_workers=self.num_workers, prefetch_factor=4)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_ds, batch_size=self.batch_size,
+                          shuffle=False, num_workers=self.num_workers, prefetch_factor=1)
+
+
 class AIA_GOESDataModule(LightningDataModule):
     """
     PyTorch Lightning DataModule for managing AIA-GOES datasets.

forecasting/inference/ablation_inference.py

@@ -0,0 +1,189 @@
+"""
+Ablation Inference Script — Gaussian Noise Channel Masking
+==========================================================
+
+Runs inference with the pretrained model while applying Gaussian noise to
+specific AIA wavelength channels, one condition at a time. This lets you
+measure how much each channel (or combination of channels) contributes to
+forecast skill.
+
+Each ablation condition produces its own output CSV so results can be
+compared directly against the clean baseline.
+"""
+
+import argparse
+import re
+import sys
+import gc
+import pandas as pd
+import torch
+import numpy as np
+from pathlib import Path
+import yaml
+from tqdm import tqdm
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent.absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from forecasting.data_loaders.SDOAIA_dataloader import AIA_GOESDataset, NoisyAIA_GOESDataset
+import forecasting.models as models
+# Reuse inference helpers from the main inference script
+from forecasting.inference.inference import load_model_from_config, evaluate_model_on_dataset
+
+
+def build_dataset(config_data, noisy_wavelengths, noise_std):
+    """
+    Build an AIA_GOESDataset (or NoisyAIA_GOESDataset if wavelengths are given).
+
+    Parameters
+    ----------
+    config_data : dict
+    noisy_wavelengths : list of int
+        Empty list → clean baseline (plain AIA_GOESDataset).
+    noise_std : float or dict
+
+    Returns
+    -------
+    dataset : torch.utils.data.Dataset
+    """
+    aia_dir = config_data['data']['aia_dir']
+    sxr_dir = config_data['data']['sxr_dir']
+    wavelengths = config_data['wavelengths']
+    prediction_only = config_data.get('prediction_only', 'false').lower() == 'true'
+
+    common_kwargs = dict(
+        aia_dir=aia_dir,
+        sxr_dir=sxr_dir,
+        wavelengths=wavelengths,
+        only_prediction=prediction_only,
+    )
+
+    if noisy_wavelengths:
+        return NoisyAIA_GOESDataset(
+            **common_kwargs,
+            noisy_wavelengths=noisy_wavelengths,
+            noise_std=noise_std,
+        )
+    else:
+        return AIA_GOESDataset(**common_kwargs)
+
+
+def run_condition(model, dataset, label, config_data, model_params, output_dir):
+    """
+    Run inference for a single ablation condition and save results to CSV.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+    dataset : torch.utils.data.Dataset
+    label : str
+        Human-readable name for this condition (used as filename stem).
+    config_data : dict
+    model_params : dict
+    output_dir : Path
+    """
+    times = dataset.samples
+    total_samples = len(times)
+    batch_size = model_params.get('batch_size', 10)
+    input_size = model_params.get('input_size', 512)
+    patch_size = model_params.get('patch_size', 16)
+    save_weights = not model_params.get('no_weights', True)
+    save_flux = not model_params.get('no_flux', True)
+    prediction_only = config_data.get('prediction_only', 'false').lower() == 'true'
+
+    print(f"\n{'='*60}")
+    print(f"  Condition: {label}  ({total_samples} samples)")
+    print(f"{'='*60}")
+
+    timestamps, predictions, ground_truths = [], [], []
+
+    pbar = tqdm(
+        evaluate_model_on_dataset(
+            model, dataset, batch_size, times, config_data,
+            save_weights, input_size, patch_size, save_flux,
+        ),
+        total=total_samples,
+        desc=label,
+        unit="sample",
+        ncols=100,
+    )
+
+    for prediction, sxr, _weight, _flux, idx in pbar:
+        predictions.append(float(prediction.item() if hasattr(prediction, 'item') else prediction))
+        ground_truths.append(0.0 if prediction_only else
+                             float(sxr.item() if hasattr(sxr, 'item') else sxr))
+        timestamps.append(str(times[idx]))
+
+    output_path = output_dir / f"{label}.csv"
+    pd.DataFrame({
+        'timestamp': timestamps,
+        'predictions': predictions,
+        'groundtruth': ground_truths,
+    }).to_csv(output_path, index=False)
+
+    print(f"  Saved → {output_path}")
+
+    # Free dataset memory between conditions
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def resolve_config_variables(config_dict):
+    """Recursively resolve ${variable} references within config."""
+    variables = {k: v for k, v in config_dict.items()
+                 if isinstance(v, str) and not v.startswith('${')}
+
+    def substitute(obj):
+        if isinstance(obj, dict):
+            return {k: substitute(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [substitute(item) for item in obj]
+        elif isinstance(obj, str):
+            for match in re.finditer(r'\$\{([^}]+)\}', obj):
+                var = match.group(1)
+                if var in variables:
+                    obj = obj.replace(f'${{{var}}}', variables[var])
+            return obj
+        return obj
+
+    return substitute(config_dict)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Ablation inference with Gaussian noise masking.")
+    parser.add_argument('-config', type=str, required=True,
+                        help='Path to the ablation inference YAML config.')
+    args = parser.parse_args()
+
+    with open(args.config, 'r') as f:
+        config_data = yaml.load(f, Loader=yaml.SafeLoader)
+    config_data = resolve_config_variables(config_data)
+    sys.modules['models'] = models
+
+    model_params = config_data.get('model_params', {})
+    ablation_cfg = config_data.get('ablation', {})
+    noise_std = ablation_cfg.get('noise_std', 1.0)
+    conditions = ablation_cfg.get('conditions', [])
+
+    if not conditions:
+        raise ValueError("No ablation conditions defined in config under 'ablation.conditions'.")
+
+    output_dir = Path(config_data['output_dir'])
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Loading model...")
+    model = load_model_from_config(config_data)
+
+    for condition in conditions:
+        label = condition.get('label', 'unnamed')
+        noisy_wavelengths = condition.get('wavelengths', [])
+
+        dataset = build_dataset(config_data, noisy_wavelengths, noise_std)
+        run_condition(model, dataset, label, config_data, model_params, output_dir)
+
+    print(f"\nAll conditions complete. Results saved to: {output_dir}")
+
+
+if __name__ == '__main__':
+    main()

forecasting/inference/ablation_inference_config.yaml

@@ -0,0 +1,62 @@
+# =============================================================================
+# FOXES Ablation Inference Configuration
+# =============================================================================
+# Used by ablation_inference.py to run Gaussian noise channel-masking ablation.
+#
+# Usage:
+#   python ablation_inference.py -config ablation_inference_config.yaml
+#
+# Variables
+# ---------
+# Define top-level string keys and reference them anywhere with ${key}.
+
+base_dir:   "/Volumes/T9/FOXES_Data"
+checkpoint: "/Users/griffingoodwin/Downloads/FOXES_Model_Checkpoint.ckpt"
+
+model: "ViTLocal"
+wavelengths: [94, 131, 171, 193, 211, 304, 335]
+prediction_only: "false"
+
+data:
+  aia_dir:         "${base_dir}/AIA_processed/val"
+  sxr_dir:         "${base_dir}/SXR_processed/val"
+  sxr_norm_path:   "${base_dir}/SXR_processed/normalized_sxr.npy"
+  checkpoint_path: "${checkpoint}"
+
+# Output directory — each condition saves <label>.csv here
+output_dir: "${base_dir}/inference/ablation"
+
+model_params:
+  input_size: 512
+  patch_size: 8
+  batch_size: 10
+  no_weights: true   # skip attention saving for speed
+  no_flux:    true
+
+# -----------------------------------------------------------------------------
+# Ablation conditions
+# -----------------------------------------------------------------------------
+# noise_std: standard deviation of Gaussian noise applied in normalized tensor space.
+# Each condition names a wavelength subset to corrupt; empty list = clean baseline.
+# You can also pass a per-wavelength dict, e.g. {171: 0.5, 193: 2.0}
+ablation:
+  noise_std: 1.0
+  conditions:
+    - label: baseline
+      wavelengths: []
+    - label: ablate_94
+      wavelengths: [94]
+    - label: ablate_131
+      wavelengths: [131]
+    - label: ablate_171
+      wavelengths: [171]
+    - label: ablate_193
+      wavelengths: [193]
+    - label: ablate_211
+      wavelengths: [211]
+    - label: ablate_304
+      wavelengths: [304]
+    - label: ablate_335
+      wavelengths: [335]
+    - label: ablate_all
+      wavelengths: [94, 131, 171, 193, 211, 304, 335]

pipeline_config.yaml

@@ -7,6 +7,14 @@
 #   python run_pipeline.py --config pipeline_config.yaml --steps all
 #   python run_pipeline.py --config pipeline_config.yaml --steps train,inference,flare_analysis
 #   python run_pipeline.py --list
+#
+# Variables
+# ---------
+# Define top-level string keys and reference them anywhere with ${key}.
+# Change base_dir or checkpoint once and every path updates automatically.
+
+base_dir:   "/Volumes/T9/FOXES_Data"
+checkpoint: "/Users/griffingoodwin/Downloads/FOXES_Model_Checkpoint.ckpt"
 
 # -----------------------------------------------------------------------------
 # HuggingFace download  (step: hf_download)
@@ -14,7 +22,7 @@
 # Use this instead of download_aia + download_sxr + preprocess + split.
 # -----------------------------------------------------------------------------
 hf_download:
-  config: "download/hf_download_config.yaml"   # points to hf_download_config.yaml
+  config: "download/hf_download_config.yaml"
 
 # -----------------------------------------------------------------------------
 # Shared date range (used by download_aia and download_sxr)
@@ -26,36 +34,36 @@ end_date:   "2014-07-08 00:00:00"
 # AIA download  (step: download_aia)
 # -----------------------------------------------------------------------------
 aia:
-  download_dir: "/Volumes/T9/Data_FOXES/AIA_raw"
+  download_dir: "${base_dir}/AIA_raw"
   email: ""        # Must be registered at http://jsoc.stanford.edu
-  cadence: 1                    # Minutes between frames
+  cadence: 1       # Minutes between frames
 
 # -----------------------------------------------------------------------------
 # SXR download  (step: download_sxr)
 # -----------------------------------------------------------------------------
 sxr:
-  save_dir: "/Volumes/T9/Data_FOXES/SXR_raw"
+  save_dir: "${base_dir}/SXR_raw"
 
 # -----------------------------------------------------------------------------
 # Preprocessing  (step: preprocess)
 # -----------------------------------------------------------------------------
 preprocess:
-  config: "data/pipeline_config.yaml"   # PipelineConfig for process_data_pipeline.py
+  config: "data/pipeline_config.yaml"
 
 # -----------------------------------------------------------------------------
 # SXR normalization  (step: normalize)
 # -----------------------------------------------------------------------------
 normalize:
-  sxr_dir:     "/Volumes/T9/Data_FOXES/SXR_processed/train"
-  output_path: "/Volumes/T9/Data_FOXES/SXR_processed/normalized_sxr.npy"
+  sxr_dir:     "${base_dir}/SXR_processed/train"
+  output_path: "${base_dir}/SXR_processed/normalized_sxr.npy"
 
 # -----------------------------------------------------------------------------
 # Train/val/test split  (step: split)
 # Runs split_data.py once for AIA and once for SXR.
 # -----------------------------------------------------------------------------
 split:
-  aia_input_dir: "/Volumes/T9/Data_FOXES/AIA_processed"  # splits into AIA_processed/train|val|test
-  sxr_input_dir: "/Volumes/T9/Data_FOXES/SXR_processed"  # splits into SXR_processed/train|val|test
+  aia_input_dir: "${base_dir}/AIA_processed"
+  sxr_input_dir: "${base_dir}/SXR_processed"
   train_start: "2014-07-01"
   train_end:   "2014-07-05"
   val_start:   "2014-07-06"
@@ -68,14 +76,14 @@ split:
 # -----------------------------------------------------------------------------
 train:
   config: "forecasting/training/train_config.yaml"
-  overrides:                          # Any key from train_config.yaml can go here
-    base_data_dir: "/Volumes/T9/Data_FOXES"
-    base_checkpoint_dir: "/Volumes/T9/Data_FOXES"
+  overrides:
+    base_data_dir:       "${base_dir}"
+    base_checkpoint_dir: "${base_dir}"
     epochs: 150
     batch_size: 6
     wandb:
       run_name: "pipeline-run"
-      entity: jayantbiradar619-university-of-arizona # Use your exact W&B username
+      entity: jayantbiradar619-university-of-arizona
       project: Paper
       job_type: training
       tags:
@@ -89,32 +97,53 @@ train:
 # -----------------------------------------------------------------------------
 inference:
   config: "forecasting/inference/local_config.yaml"
-  overrides:                          # Any key from local_config.yaml can go here
+  overrides:
     data:
-      aia_dir:         "/Volumes/T9/Data_FOXES/AIA_processed"
-      sxr_dir:         "/Volumes/T9/Data_FOXES/SXR_processed"
-      sxr_norm_path:   "/Volumes/T9/Data_FOXES/SXR_processed/normalized_sxr.npy"
-      checkpoint_path: "/Users/griffingoodwin/Downloads/FOXES_Model_Checkpoint.ckpt"  # update to actual checkpoint
-    output_path: "/Volumes/T9/Data_FOXES/inference/predictions.csv"
+      aia_dir:         "${base_dir}/AIA_processed"
+      sxr_dir:         "${base_dir}/SXR_processed"
+      sxr_norm_path:   "${base_dir}/SXR_processed/normalized_sxr.npy"
+      checkpoint_path: "${checkpoint}"
+    output_path: "${base_dir}/inference/predictions.csv"
     prediction_only: "false"
     paths:
-      data_dir:        "/Volumes/T9/Data_FOXES"
-      predictions_csv: "/Volumes/T9/Data_FOXES/inference/predictions.csv"
-      aia_path:        "/Volumes/T9/Data_FOXES/AIA_processed/test"
+      data_dir:        "${base_dir}"
+      predictions_csv: "${base_dir}/inference/predictions.csv"
+      aia_path:        "${base_dir}/AIA_processed/test"
+
+# -----------------------------------------------------------------------------
+# Ablation study  (step: ablation)
+# -----------------------------------------------------------------------------
+ablation:
+  config: "forecasting/inference/ablation_inference_config.yaml"
+  overrides:
+    data:
+      aia_dir:         "${base_dir}/AIA_processed/test"
+      sxr_dir:         "${base_dir}/SXR_processed/test"
+      sxr_norm_path:   "${base_dir}/SXR_processed/normalized_sxr.npy"
+      checkpoint_path: "${checkpoint}"
+    output_dir: "${base_dir}/inference/ablation"
+
+# -----------------------------------------------------------------------------
+# Spatial performance  (step: spatial_performance)
+# -----------------------------------------------------------------------------
+spatial_performance:
+  flux_dir:        "${base_dir}/flux"
+  predictions_csv: "${base_dir}/inference/predictions.csv"
+  out_dir:         "${base_dir}/inference/spatial_performance"
 
 # -----------------------------------------------------------------------------
 # Evaluation  (step: evaluate)
 # -----------------------------------------------------------------------------
 evaluate:
   config: "forecasting/inference/evaluation_config.yaml"
-  overrides:                          # Any key from evaluation_config.yaml can go here
+  overrides:
     model_predictions:
-      main_model_csv: "/Volumes/T9/Data_FOXES/inference/predictions.csv"
+      main_model_csv: "${base_dir}/inference/predictions.csv"
     data:
-      aia_dir:     "/Volumes/T9/Data_FOXES/AIA_processed/val"
-      weight_path: "/Volumes/T9/Data_FOXES/inference/weights"
+      aia_dir:     "${base_dir}/AIA_processed/val"
+      weight_path: "${base_dir}/inference/weights"
     evaluation:
-      output_dir:  "/Volumes/T9/Data_FOXES/inference/evaluation"
+      output_dir:  "${base_dir}/inference/evaluation"
     time_range:
       start_time:       "2023-01-01T00:00:00"
       end_time:         "2023-12-31T23:59:59"

run_pipeline.py

@@ -23,6 +23,7 @@ Usage:
 
 import argparse
 import logging
+import re
 import subprocess
 import sys
 import time
@@ -47,6 +48,33 @@ log = logging.getLogger(__name__)
 # Config helpers
 # ---------------------------------------------------------------------------
 
+def resolve_variables(cfg: dict) -> dict:
+    """
+    Resolve ${variable} placeholders in a config dict.
+
+    Top-level string keys whose values don't themselves contain placeholders
+    are treated as variables. Substitution is applied recursively to all
+    string values in the config.
+    """
+    variables = {k: v for k, v in cfg.items()
+                 if isinstance(v, str) and '${' not in v}
+
+    def substitute(obj):
+        if isinstance(obj, dict):
+            return {k: substitute(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [substitute(item) for item in obj]
+        elif isinstance(obj, str):
+            for match in re.finditer(r'\$\{([^}]+)\}', obj):
+                var = match.group(1)
+                if var in variables:
+                    obj = obj.replace(f'${{{var}}}', variables[var])
+            return obj
+        return obj
+
+    return substitute(cfg)
+
+
 def deep_merge(base: dict, overrides: dict) -> dict:
     """Recursively merge overrides into base, modifying base in-place."""
     for key, val in overrides.items():
@@ -88,6 +116,8 @@ STEP_ORDER = [
     "inference",
     "evaluate",
     "flare_analysis",
+    "ablation",
+    "spatial_performance",
 ]
 
 STEP_INFO = {
@@ -135,6 +165,14 @@ STEP_INFO = {
         "description": "Detect, track, and match flares; generate plots/movies",
         "script": ROOT / "forecasting" / "inference" / "flare_analysis.py",
     },
+    "ablation": {
+        "description": "Run Gaussian noise channel-masking ablation study on pretrained model",
+        "script": ROOT / "forecasting" / "inference" / "ablation_inference.py",
+    },
+    "spatial_performance": {
+        "description": "Generate flux-weighted spatial error heatmap on the solar disk",
+        "script": ROOT / "analysis" / "spatial_performance.py",
+    },
 }
 
 
@@ -268,6 +306,26 @@ def build_commands(step: str, cfg: dict, force: bool) -> list[list[str]] | None:
             config_path = str(write_merged_config(config_path, inf["overrides"], "inference_config"))
         return [base + ["--config", config_path]]
 
+    if step == "ablation":
+        if not require(["config"], "ablation"):
+            return None
+        abl = cfg["ablation"]
+        config_path = abl["config"]
+        if abl.get("overrides"):
+            config_path = str(write_merged_config(config_path, abl["overrides"], "ablation_config"))
+        return [base + ["-config", config_path]]
+
+    if step == "spatial_performance":
+        sp = cfg.get("spatial_performance", {})
+        cmd = base[:]
+        if sp.get("flux_dir"):
+            cmd += ["--flux_dir", sp["flux_dir"]]
+        if sp.get("predictions_csv"):
+            cmd += ["--predictions_csv", sp["predictions_csv"]]
+        if sp.get("out_dir"):
+            cmd += ["--out_dir", sp["out_dir"]]
+        return [cmd]
+
     return [base]
 
 
@@ -338,6 +396,7 @@ def main():
 
     with open(args.config, "r") as f:
         cfg = yaml.safe_load(f)
+    cfg = resolve_variables(cfg)
 
     # Resolve step list
     if args.steps.strip().lower() == "all":