refactiring rest of code base and adding checkpoints

.gitignore

@@ -172,3 +172,4 @@ analysis/performance_heatmap_all.png
 download/download_flares.py
 download/test_download_sdo.py
 forecasting/models/test_vit_mask.py
+forecasting/griffin_config.yaml

README.md

@@ -46,20 +46,18 @@ The solar soft X-ray (SXR) irradiance is a long-standing proxy of solar activity
 
 ## Repository Structure
 
-This repository is intentionally scoped to **getting data, running inference,
-and evaluating** a trained FOXES model — it does not include the training
-code used to produce the released checkpoint.
+This repository covers the full loop: getting data, training, running
+inference, and evaluating a FOXES model.
 
 ```text
 FOXES
 ├── download
-│   ├── hugging_face_data_download.py # Recommended: stream FOXES data from HF Hub straight to .npy
-│   ├── parquet_to_npy.py             # Convert already-downloaded HF parquet files to .npy
+│   ├── hugging_face_data_download.py # Recommended: HF Hub -> .npy (streamed, or from local parquet)
 │   ├── hf_download_config.yaml       # Config for hugging_face_data_download.py
 │   ├── download_sdo.py               # Advanced: raw AIA download from JSOC (needs data/build_dataset.py after)
 │   ├── sdo_download_config.yaml      # Config for download_sdo.py
-│   ├── sxr_downloader.py             # Advanced: raw GOES SXR download via SunPy Fido (needs data/build_dataset.py after)
-│   └── sxr_download_config.yaml      # Config for sxr_downloader.py
+│   ├── download_sxr.py               # Advanced: raw GOES SXR download via SunPy Fido (needs data/build_dataset.py after)
+│   └── sxr_download_config.yaml      # Config for download_sxr.py
 ├── data
 │   ├── build_dataset.py         # Runs the full raw -> processed pipeline below in one command
 │   ├── build_dataset_config.yaml # Config for build_dataset.py
@@ -67,14 +65,19 @@ FOXES
 │   ├── convert_aia.py           # Raw AIA FITS -> paired 512x512 .npy stacks (itipy)
 │   ├── combine_sxr.py           # Combine raw multi-satellite GOES files into per-satellite CSVs
 │   ├── align_aia_sxr.py         # Match AIA timestamps to GOES CSVs -> per-timestamp SXR .npy
+│   ├── split_train_val_test.py  # Split processed AIA/SXR into train/val/test (training only)
 │   └── sxr_normalization.py     # Compute log-space mean/std over SXR .npy files for training
 ├── forecasting
-│   ├── dataset.py            # AIA_GOESDataset: loads paired AIA + SXR .npy files
+│   ├── dataset.py            # AIAGOESDataset / AIAGOESDataModule: loads paired AIA + SXR .npy files
 │   ├── model.py               # ViTLocal: Vision Transformer with patch flux heads
 │   ├── inference.py           # Run a checkpoint over a folder of data; writes predictions.csv
 │   ├── inference_config.yaml  # Config for inference.py
 │   ├── evaluation.py          # Compute metrics and generate evaluation plots
 │   └── evaluation_config.yaml # Config for evaluation.py
+├── training
+│   ├── train.py                # Train ViTLocal with PyTorch Lightning + Weights & Biases logging
+│   ├── train_config.yaml       # Config for train.py
+│   └── callbacks.py            # W&B callbacks: SXR pred-vs-true plots, attention map visualization
 └── requirements.txt            # Python dependencies
 ```
 
@@ -107,9 +110,10 @@ conda activate foxes
 
 ## Running the Model
 
-FOXES is run in three steps: **get data**, **inference** (run a checkpoint over
-your data), and **evaluation** (score the predictions and generate plots). All
-three are driven by a YAML config — edit the config, then run the script.
+FOXES is run in four steps: **get data**, **inference** (run a checkpoint over
+your data), **evaluation** (score the predictions and generate plots), and
+optionally **training** your own model. All are driven by a YAML config — edit
+the config, then run the script.
 
 ### 0) Get data
 
@@ -124,7 +128,9 @@ Edit `download/hf_download_config.yaml` first to set `aia_dir`/`sxr_dir`, which
 splits to pull, and whether to subsample.
 
 If you've already downloaded the HF parquet files locally instead of
-streaming, convert them the same way with `download/parquet_to_npy.py`.
+streaming, set `local_parquet_dir` in that same config to the root folder
+containing your per-split subdirs (`train/`, `validation/` or `val/`, `test/`)
+— streaming and the HF Hub login are skipped entirely in that case.
 
 **Advanced:** to acquire and process raw data yourself instead, edit and run
 the two download configs, then the one dataset-build config, in order:
@@ -134,18 +140,22 @@ the two download configs, then the one dataset-build config, in order:
 python download/download_sdo.py --config download/sdo_download_config.yaml
 
 # 2) Raw GOES XRS data via SunPy Fido
-python download/sxr_downloader.py --config download/sxr_download_config.yaml
+python download/download_sxr.py --config download/sxr_download_config.yaml
 
 # 3) Clean + convert AIA, combine + align SXR -> paired .npy (see data/build_dataset_config.yaml)
-python data/build_dataset.py -config data/build_dataset_config.yaml
+python data/build_dataset.py --config data/build_dataset_config.yaml
 ```
 
 `data/build_dataset.py` runs the full raw-to-processed pipeline in one command
 (clean AIA → convert AIA → combine GOES → align AIA/SXR); each step can be
-skipped via the `steps:` block in its config if you've already run it. It can
-also optionally compute SXR normalization stats for training — set
-`sxr_normalization.compute: true` in the config (not needed for inference
-against a released checkpoint).
+skipped via the `steps:` block in its config if you've already run it.
+Inference/evaluation just need the flat output of that — training needs two
+more things, both off by default and only relevant if you're training:
+- `steps.split: true` — splits `aia.processed_dir`/`output.sxr_dir` into
+  `train/`/`val/`/`test/` subfolders (date ranges or a month-based default;
+  see the `split:` block in the config).
+- `sxr_normalization.compute: true` — computes SXR normalization stats from
+  the train split (requires `steps.split` to have run first).
 
 ### 1) Data format
 
@@ -186,7 +196,7 @@ output_path: "/path/to/predictions.csv"
 Then run:
 
 ```bash
-python forecasting/inference.py -config forecasting/inference_config.yaml
+python forecasting/inference.py --config forecasting/inference_config.yaml
 ```
 
 This writes `output_path` (a CSV of timestamp/prediction/groundtruth). Per-patch
@@ -201,11 +211,46 @@ Edit `forecasting/evaluation_config.yaml` to point at the predictions
 CSV and data directories, then run:
 
 ```bash
-python forecasting/evaluation.py -config forecasting/evaluation_config.yaml
+python forecasting/evaluation.py --config forecasting/evaluation_config.yaml
 ```
 
 This computes metrics (MSE, MAE, R²) and generates plots under `evaluation.output_dir`.
 
+### 4) Train your own model (optional)
+
+Unlike inference, training expects `aia_dir`/`sxr_dir` each to have `train/`,
+`val/`, and `test/` subfolders of paired `.npy` files — exactly what
+`data/build_dataset.py` and the Hugging Face download path produce.
+
+Edit `training/train_config.yaml`:
+
+```yaml
+base_data_dir: "/path/to/processed_data"        # holds AIA_processed/ and SXR_processed/
+base_checkpoint_dir: "/path/to/checkpoints"
+
+gpu_ids: -1        # -1 = CPU, 0 = GPU 0, [0,1] = specific GPUs, "all" = every GPU
+batch_size: 6
+epochs: 150
+
+vit_architecture:
+  mask_mode: inverted   # inverted (released model) | local | none (full/global attention)
+  local_window: 9
+
+wandb:
+  entity: ""   # your W&B username or team name
+```
+
+Then run:
+
+```bash
+python training/train.py --config training/train_config.yaml
+```
+
+Training logs to Weights & Biases (predicted-vs-true SXR plots and attention
+map visualizations each validation epoch — see `training/callbacks.py`) and
+saves the top 10 checkpoints by validation loss to `data.checkpoints_dir`,
+ready to point `forecasting/inference_config.yaml` at.
+
 ---
 
 ## Citation

data/align_aia_sxr.py

@@ -171,7 +171,10 @@ def align_aia_sxr(goes_data_dir, aia_processed_dir, output_sxr_dir, aia_missing_
 
     # Get target timestamps from AIA files
     print(f"\nFinding target timestamps from AIA files in: {aia_processed_dir}")
-    aia_files = sorted(glob.glob(f"{aia_processed_dir}/*.npy", recursive=True))
+    # Skip macOS AppleDouble sidecar files (e.g. "._2023-08-01T00:00:00.npy")
+    # that external/non-native filesystems litter alongside real files.
+    aia_files = sorted(f for f in glob.glob(f"{aia_processed_dir}/*.npy", recursive=True)
+                       if not os.path.basename(f).startswith('._'))
     aia_files_split = [file.split('/')[-1].split('.')[0] for file in aia_files]
     common_timestamps = [
         datetime.fromisoformat(date_str).strftime('%Y-%m-%dT%H:%M:%S')

data/build_dataset.py

@@ -1,15 +1,16 @@
 """
 End-to-end raw -> processed dataset builder for FOXES.
 
-Runs, in order: clean_aia -> convert_aia -> combine_sxr -> align -> (optional)
-sxr_normalization. Each step can be skipped via the `steps` section of the
-config if you already ran it. Everything is driven by one YAML config — see
-build_dataset_config.yaml for the fields.
+Runs, in order: clean_aia -> convert_aia -> combine_sxr -> align -> split ->
+(optional) sxr_normalization. Each step can be skipped via the `steps` section
+of the config if you already ran it. Everything is driven by one YAML config
+— see build_dataset_config.yaml for the fields.
 
 Usage:
-    python data/build_dataset.py -config data/build_dataset_config.yaml
+    python data/build_dataset.py --config data/build_dataset_config.yaml
 """
 import argparse
+import os
 
 import numpy as np
 import yaml
@@ -18,6 +19,7 @@ from clean_aia import clean_aia_data
 from convert_aia import process_aia_to_npy
 from combine_sxr import SXRDataProcessor
 from align_aia_sxr import align_aia_sxr
+from split_train_val_test import split_train_val_test
 from sxr_normalization import compute_sxr_norm
 
 
@@ -28,7 +30,7 @@ def load_config(config_path):
 
 def main():
     parser = argparse.ArgumentParser(description='Build a FOXES-ready dataset from raw AIA/GOES data.')
-    parser.add_argument('-config', type=str, default='data/build_dataset_config.yaml',
+    parser.add_argument('--config', type=str, default='data/build_dataset_config.yaml',
                         help='Path to build_dataset config YAML file')
     args = parser.parse_args()
 
@@ -40,30 +42,31 @@ def main():
     output = config['output']
     processing = config.get('processing', {})
     steps = config.get('steps', {})
+    split = config.get('split', {})
     norm = config.get('sxr_normalization', {})
 
     print("=== FOXES Dataset Build ===")
 
     if steps.get('clean_aia', True):
-        print("\n--- Step 1/4: Clean AIA FITS files ---")
+        print("\n--- Step 1/5: Clean AIA FITS files ---")
         clean_aia_data(aia['raw_dir'], aia['bad_files_dir'], wavelengths)
     else:
-        print("\n--- Step 1/4: Clean AIA FITS files (skipped) ---")
+        print("\n--- Step 1/5: Clean AIA FITS files (skipped) ---")
 
     if steps.get('convert_aia', True):
-        print("\n--- Step 2/4: Convert AIA FITS -> .npy ---")
+        print("\n--- Step 2/5: Convert AIA FITS -> .npy ---")
         process_aia_to_npy(aia['raw_dir'], aia['processed_dir'], wavelengths)
     else:
-        print("\n--- Step 2/4: Convert AIA FITS -> .npy (skipped) ---")
+        print("\n--- Step 2/5: Convert AIA FITS -> .npy (skipped) ---")
 
     if steps.get('combine_sxr', True):
-        print("\n--- Step 3/4: Combine raw GOES satellite files ---")
+        print("\n--- Step 3/5: Combine raw GOES satellite files ---")
         SXRDataProcessor(data_dir=sxr['raw_dir'], output_dir=sxr['combined_dir']).combine_goes_data()
     else:
-        print("\n--- Step 3/4: Combine raw GOES satellite files (skipped) ---")
+        print("\n--- Step 3/5: Combine raw GOES satellite files (skipped) ---")
 
     if steps.get('align', True):
-        print("\n--- Step 4/4: Align AIA timestamps with GOES SXR data ---")
+        print("\n--- Step 4/5: Align AIA timestamps with GOES SXR data ---")
         align_aia_sxr(
             goes_data_dir=sxr['combined_dir'],
             aia_processed_dir=aia['processed_dir'],
@@ -74,11 +77,22 @@ def main():
             min_batch_size=processing.get('min_batch_size', 1),
         )
     else:
-        print("\n--- Step 4/4: Align AIA timestamps with GOES SXR data (skipped) ---")
+        print("\n--- Step 4/5: Align AIA timestamps with GOES SXR data (skipped) ---")
+
+    # Off by default: only needed if you're training (inference/evaluation
+    # work fine against the flat processed_dir/sxr_dir from the align step).
+    if steps.get('split', False):
+        print("\n--- Step 5/5: Split AIA/SXR into train/val/test ---")
+        split_kwargs = dict(train_range=split.get('train_range'), val_range=split.get('val_range'),
+                            test_range=split.get('test_range'), copy_files=split.get('copy_files', False))
+        split_train_val_test(aia['processed_dir'], aia['processed_dir'], **split_kwargs)
+        split_train_val_test(output['sxr_dir'], output['sxr_dir'], **split_kwargs)
+    else:
+        print("\n--- Step 5/5: Split AIA/SXR into train/val/test (skipped) ---")
 
     if norm.get('compute', False):
-        print("\n--- Optional: Compute SXR normalization stats ---")
-        sxr_norm = compute_sxr_norm(output['sxr_dir'])
+        print("\n--- Optional: Compute SXR normalization stats (from the train split) ---")
+        sxr_norm = compute_sxr_norm(os.path.join(output['sxr_dir'], 'train'))
         np.save(norm['output_path'], sxr_norm)
         print(f"Saved SXR normalization to {norm['output_path']}")
 

data/build_dataset_config.yaml

@@ -4,23 +4,23 @@
 # Turns raw AIA FITS + raw GOES XRS data into the paired .npy layout that
 # forecasting/dataset.py expects.
 #
-# Usage: python data/build_dataset.py -config data/build_dataset_config.yaml
+# Usage: python data/build_dataset.py --config data/build_dataset_config.yaml
 # =============================================================================
 
 wavelengths: [94, 131, 171, 193, 211, 304, 335]
 
 aia:
-  raw_dir:        "/path/to/AIA_raw"       # raw FITS from download/download_sdo.py
-  bad_files_dir:  "/path/to/AIA_bad"       # bad-timestamp FITS get moved here
-  processed_dir:  "/path/to/AIA_processed" # paired 512x512 .npy stacks
+  raw_dir:        "/Volumes/T9/testing_foxes/AIA_raw"       # raw FITS from download/download_sdo.py
+  bad_files_dir:  "/Volumes/T9/testing_foxes/SDO/AIA_bad"       # bad-timestamp FITS get moved here
+  processed_dir:  "/Volumes/T9/testing_foxes/AIA_processed" # paired 512x512 .npy stacks
 
 sxr:
-  raw_dir:      "/path/to/GOES_raw"        # raw netCDF from download/sxr_downloader.py
-  combined_dir: "/path/to/GOES_combined"   # one interpolated CSV per satellite
+  raw_dir:      "/Volumes/T9/testing_foxes/GOES_raw"        # raw netCDF from download/download_sxr.py
+  combined_dir: "/Volumes/T9/testing_foxes/GOES_combined"   # one interpolated CSV per satellite
 
 output:
-  sxr_dir:         "/path/to/SXR_processed" # per-timestamp xrsb_flux .npy, aligned to AIA
-  aia_missing_dir: "/path/to/AIA_missing"   # AIA files with no matching SXR value
+  sxr_dir:         "/Volumes/T9/testing_foxes/SXR_processed" # per-timestamp xrsb_flux .npy, aligned to AIA
+  aia_missing_dir: "/Volumes/T9/testing_foxes/AIA_missing"   # AIA files with no matching SXR value
 
 processing:
   max_processes: null    # cap worker processes; null = use all CPUs
@@ -30,13 +30,26 @@ processing:
 # Which build steps to run, in order. Set any to false to skip it (e.g. if
 # you already ran that step and are re-running after fixing a later one).
 steps:
-  clean_aia:   true   # drop AIA FITS files with a bad DATE-OBS timestamp
-  convert_aia: true   # raw AIA FITS -> paired .npy stacks
-  combine_sxr: true   # combine raw multi-satellite GOES files into per-satellite CSVs
-  align:       true   # match AIA timestamps to GOES CSVs -> per-timestamp SXR .npy
+  clean_aia:   true    # drop AIA FITS files with a bad DATE-OBS timestamp
+  convert_aia: true    # raw AIA FITS -> paired .npy stacks
+  combine_sxr: true    # combine raw multi-satellite GOES files into per-satellite CSVs
+  align:       true    # match AIA timestamps to GOES CSVs -> per-timestamp SXR .npy
+  split:       true   # split aia.processed_dir / output.sxr_dir into train/val/test
+
+# Only needed if you're training (training/train_config.yaml expects
+# aia_dir/sxr_dir to each have train/val/test subfolders) — not needed for
+# inference or evaluation against a released checkpoint. Splits both
+# aia.processed_dir and output.sxr_dir in place, using the same date ranges
+# so a given timestamp lands in the same split for both.
+split:
+  train_range: ["2023-08-01","2023-08-15"]   # e.g. ["2012-01-01", "2022-12-31"]; null with the others -> month-based default
+  val_range:   ["2023-08-16","2023-08-22"]   # e.g. ["2023-01-01", "2023-06-30"]
+  test_range:  ["2023-08-23","2023-08-31"]   # e.g. ["2023-08-01", "2023-08-31"]
+  copy_files:  false  # copy instead of move
 
 # Optional — only needed if you're training, not for running inference
-# against a released checkpoint.
+# against a released checkpoint. Requires steps.split to have run (computes
+# stats from the train split only).
 sxr_normalization:
-  compute:     false
-  output_path: "/path/to/normalized_sxr.npy"
+  compute:     true
+  output_path: "/Volumes/T9/testing_foxes/normalized_sxr.npy"

data/clean_aia.py

@@ -1,6 +1,11 @@
 """
-Remove AIA FITS files whose filename timestamp doesn't match their DATE-OBS
-header (more than 60s apart), which indicates a mis-tagged or corrupt download.
+Remove AIA FITS files that are unusable: filename timestamp doesn't match the
+DATE-OBS header (more than 60s apart, indicating a mis-tagged/corrupt download),
+the image itself is blank (all-zero/constant/NaN), or QUALITY is nonzero (AIA's
+own bitmask flagging eclipse/calibration/off-point/etc. frames). download_sdo.py
+only warns and falls back to the closest available frame when no quality-0
+frame exists nearby, so a bad-quality frame can still make it to disk — this
+is the backstop that catches it.
 
 Called by build_dataset.py — not meant to be run standalone.
 """
@@ -9,36 +14,51 @@ import os
 import shutil
 from multiprocessing import Pool
 
+import numpy as np
 import pandas as pd
 from tqdm import tqdm
 
 # itipy needs these aliases done manually on newer Python versions.
-collections.Iterable = collections.abc.Iterable
-collections.Mapping = collections.abc.Mapping
-collections.MutableSet = collections.abc.MutableSet
-collections.MutableMapping = collections.abc.MutableMapping
+collections.Iterable = collections.abc.Iterable  # type: ignore[attr-defined]
+collections.Mapping = collections.abc.Mapping  # type: ignore[attr-defined]
+collections.MutableSet = collections.abc.MutableSet  # type: ignore[attr-defined]
+collections.MutableMapping = collections.abc.MutableMapping  # type: ignore[attr-defined]
 from itipy.data.dataset import get_intersecting_files
 from astropy.io import fits
 
 
+def _is_blank(data):
+    """True if the image has no usable signal (all-zero, constant, or all-NaN)."""
+    if data is None:
+        return True
+    finite = data[np.isfinite(data)]
+    if finite.size == 0:
+        return True
+    return bool(np.nanmax(finite) == np.nanmin(finite))
+
+
 def process_fits_file(file_path):
     try:
         with fits.open(file_path) as hdu:
-            header = hdu[1].header
+            header = hdu[1].header  # type: ignore[union-attr]
             date_obs = pd.to_datetime(header['DATE-OBS'])
             # Ensure timezone-naive datetime
             if date_obs.tz is not None:
                 date_obs = date_obs.tz_localize(None)
             wavelength = header['WAVELNTH']
             filename = pd.to_datetime(os.path.basename(file_path).split('.')[0])
-            return {'DATE-OBS': date_obs, 'WAVELNTH': wavelength, 'FILENAME': filename}
+            blank = _is_blank(hdu[1].data)  # type: ignore[union-attr]
+            quality = pd.to_numeric(pd.Series([header.get('QUALITY')]), errors='coerce').iloc[0]
+            bad_quality = bool(pd.notna(quality) and quality != 0)
+            return {'DATE-OBS': date_obs, 'WAVELNTH': wavelength, 'FILENAME': filename,
+                    'BLANK': blank, 'BAD_QUALITY': bad_quality}
     except Exception as e:
         print(f"Error processing {file_path}: {e}")
         return None
 
 
 def clean_aia_data(input_folder, bad_files_dir, wavelengths):
-    """Move AIA FITS files with bad DATE-OBS timestamps out of input_folder."""
+    """Move AIA FITS files with bad DATE-OBS timestamps or blank images out of input_folder."""
     aia_files = get_intersecting_files(input_folder, wavelengths)
     file_list = aia_files[0]  # List of FITS file paths
 
@@ -48,6 +68,10 @@ def clean_aia_data(input_folder, bad_files_dir, wavelengths):
     # Filter out None results (in case of failed files)
     results = [r for r in results if r is not None]
 
+    if not results:
+        print("No readable AIA files found — nothing to clean.")
+        return
+
     # Convert to DataFrame
     aia_header = pd.DataFrame(results)
     aia_header['DATE-OBS'] = pd.to_datetime(aia_header['DATE-OBS'])
@@ -57,9 +81,13 @@ def clean_aia_data(input_folder, bad_files_dir, wavelengths):
         pd.to_datetime(aia_header['FILENAME']) - pd.to_datetime(aia_header['DATE-OBS'])
     ).dt.total_seconds()
 
-    # Remove rows where DATE_DIFF is greater than ±60 seconds
-    files_to_remove = aia_header[(aia_header['DATE_DIFF'] <= -60) | (aia_header['DATE_DIFF'] >= 60)]
-    print(f"{len(files_to_remove)} bad files found")
+    # Remove rows where DATE_DIFF is greater than ±60 seconds, the image is
+    # blank, or QUALITY flags a known instrumental issue
+    bad_timing = (aia_header['DATE_DIFF'] <= -60) | (aia_header['DATE_DIFF'] >= 60)
+    files_to_remove = aia_header[bad_timing | aia_header['BLANK'] | aia_header['BAD_QUALITY']]
+    print(f"{len(files_to_remove)} bad files found "
+          f"({bad_timing.sum()} bad timing, {aia_header['BLANK'].sum()} blank, "
+          f"{aia_header['BAD_QUALITY'].sum()} bad quality)")
 
     for wavelength in wavelengths:
         print(f"\nProcessing wavelength: {wavelength}")

data/combine_sxr.py

@@ -29,7 +29,7 @@ class SXRDataProcessor:
     def __init__(self, data_dir: str, output_dir: str):
         self.data_dir = Path(data_dir)
         self.output_dir = Path(output_dir)
-        self.output_dir.mkdir(exist_ok=True)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
         self.used_g13_files = []
         self.used_g14_files = []
         self.used_g15_files = []
@@ -42,14 +42,19 @@ class SXRDataProcessor:
         Combine GOES-16 and GOES-18 files and track source files used.
         Parameters
         """
-        print("🔍 Scanning for GOES data files...")
-        
-        g13_files = sorted(self.data_dir.glob("*g13*.nc"))
-        g14_files = sorted(self.data_dir.glob("*g14*.nc"))
-        g15_files = sorted(self.data_dir.glob("*g15*.nc"))
-        g16_files = sorted(self.data_dir.glob("*g16*.nc"))
-        g17_files = sorted(self.data_dir.glob("*g17*.nc"))
-        g18_files = sorted(self.data_dir.glob("*g18*.nc"))
+        print("Scanning for GOES data files...")
+
+        def _glob(pattern):
+            # Skip macOS AppleDouble sidecar files (e.g. "._sci_xrsf...nc")
+            # that external/non-native filesystems litter alongside real files.
+            return sorted(p for p in self.data_dir.glob(pattern) if not p.name.startswith('._'))
+
+        g13_files = _glob("*g13*.nc")
+        g14_files = _glob("*g14*.nc")
+        g15_files = _glob("*g15*.nc")
+        g16_files = _glob("*g16*.nc")
+        g17_files = _glob("*g17*.nc")
+        g18_files = _glob("*g18*.nc")
         
         total_files = len(g13_files) + len(g14_files) + len(g15_files) + len(g16_files) + len(g17_files) + len(g18_files)
         logging.info(
@@ -62,11 +67,10 @@ class SXRDataProcessor:
 
         def process_files(files, satellite_name, output_file, used_file_list):
             datasets = []
-            combined_meta = {}
             successful_files = 0
             failed_files = 0
 
-            print(f"🛰️  Processing {satellite_name} ({len(files)} files)...")
+            print(f"Processing {satellite_name} ({len(files)} files)...")
             
             # Progress bar for file loading
             with tqdm(files, desc=f"Loading {satellite_name}", unit="file", 
@@ -85,7 +89,7 @@ class SXRDataProcessor:
                         continue
                     finally:
                         if 'ds' in locals():
-                            ds.close()
+                            ds.close() # type: ignore[attr-defined]
 
             if not datasets:
                 print(f"No valid datasets for {satellite_name}")

data/convert_aia.py

@@ -8,10 +8,10 @@ import collections.abc
 import os
 from multiprocessing import Pool
 
-collections.Iterable = collections.abc.Iterable
-collections.Mapping = collections.abc.Mapping
-collections.MutableSet = collections.abc.MutableSet
-collections.MutableMapping = collections.abc.MutableMapping
+collections.Iterable = collections.abc.Iterable  # type: ignore[attr-defined]
+collections.Mapping = collections.abc.Mapping  # type: ignore[attr-defined]
+collections.MutableSet = collections.abc.MutableSet  # type: ignore[attr-defined]
+collections.MutableMapping = collections.abc.MutableMapping  # type: ignore[attr-defined]
 import numpy as np
 from itipy.data.dataset import StackDataset, get_intersecting_files, AIADataset
 from itipy.data.editor import BrightestPixelPatchEditor
@@ -40,7 +40,7 @@ class AIAStackDataset(StackDataset):
         ds_mapping = {94: AIADataset, 131: AIADataset, 171: AIADataset, 193: AIADataset, 211: AIADataset,
                       304: AIADataset, 335: AIADataset, 1600: AIADataset, 1700: AIADataset, 4500: AIADataset, 6173: AIADataset}
         data_sets = [ds_mapping[wl_id](files, wavelength=wl_id, resolution=resolution, ext=ext, allow_errors=allow_errors)
-                     for wl_id, files in zip(wavelengths, paths)]
+                     for wl_id, files in zip(wavelengths, paths)] # type: ignore[attr-defined]
         super().__init__(data_sets, **kwargs)
         if patch_shape is not None:
             self.addEditor(BrightestPixelPatchEditor(patch_shape))
@@ -58,11 +58,11 @@ def _init_worker(dataset, out_folder):
 
 def save_sample(i):
     try:
-        data = _aia_dataset[i]
-        file_path = os.path.join(_output_folder, _aia_dataset.getId(i)) + '.npy'
+        data = _aia_dataset[i] # type: ignore[attr-defined]
+        file_path = os.path.join(_output_folder, _aia_dataset.getId(i)) + '.npy' # type: ignore[attr-defined]
         np.save(file_path, data)
     except Exception as e:
-        print(f"Warning: Could not process sample {i} (ID: {_aia_dataset.getId(i)}): {e}")
+        print(f"Warning: Could not process sample {i} (ID: {_aia_dataset.getId(i)}): {e}") # type: ignore[attr-defined]
 
 
 def check_existing_files(base_input_folder, wavelengths, output_folder):

data/split_train_val_test.py

@@ -0,0 +1,101 @@
+"""
+Split a flat folder of timestamp-named .npy files (AIA or SXR) into
+train/val/test subfolders, so forecasting/training's AIAGOESDataModule can
+find them.
+
+Called by build_dataset.py — not meant to be run standalone.
+"""
+import os
+import shutil
+
+import pandas as pd
+
+
+def _normalize_timestamp(ts: str) -> str:
+    """Normalize timestamp strings with underscores instead of colons (cross-platform filenames)."""
+    if 'T' in ts:
+        date_part, time_part = ts.split('T', 1)
+        return f"{date_part}T{time_part.replace('_', ':')}"
+    return ts
+
+
+def _assign_split(file_time, train_range, val_range, test_range):
+    """Return 'train'/'val'/'test' for this timestamp, or None if no range matches."""
+    ranges = {'train': train_range, 'val': val_range, 'test': test_range}
+    if any(ranges.values()):
+        for split_name, rng in ranges.items():
+            if rng is None:
+                continue
+            start = pd.to_datetime(rng[0])
+            end = pd.to_datetime(rng[1]).replace(hour=23, minute=59, second=59, microsecond=999999)
+            if start <= file_time <= end:
+                return split_name
+        return None
+
+    # Default: month-based split (August held out for test, Jan-Mar for val)
+    month = file_time.month
+    if month == 8:
+        return 'test'
+    if month in (1, 2, 3):
+        return 'val'
+    return 'train'
+
+
+def split_train_val_test(input_dir, output_dir, train_range=None, val_range=None, test_range=None,
+                          copy_files=False):
+    """
+    Split .npy files in input_dir into train/val/test subfolders under output_dir,
+    based on each file's timestamp filename.
+
+    Parameters
+    ----------
+    input_dir : str
+        Flat folder of .npy files named by timestamp (e.g. from align_aia_sxr.py
+        or convert_aia.py). Can be the same path as output_dir to split in place.
+    output_dir : str
+        Destination folder; train/val/test subfolders are created under it.
+    train_range, val_range, test_range : [start, end] date strings, optional
+        Inclusive date ranges ("YYYY-MM-DD") for each split. If none are given,
+        falls back to a month-based default: August -> test, Jan-Mar -> val,
+        everything else -> train.
+    copy_files : bool
+        Copy instead of move (default: move).
+    """
+    if not os.path.isdir(input_dir):
+        raise ValueError(f"Input folder does not exist: {input_dir}")
+
+    for split_name in ("train", "val", "test"):
+        os.makedirs(os.path.join(output_dir, split_name), exist_ok=True)
+
+    files = sorted(f for f in os.listdir(input_dir) if f.endswith(".npy"))
+    print(f"Splitting {len(files)} files from {input_dir}")
+
+    moved = skipped = 0
+    for filename in files:
+        try:
+            file_time = pd.to_datetime(_normalize_timestamp(filename[:-len(".npy")]))
+        except ValueError:
+            print(f"Skipping {filename}: invalid timestamp")
+            skipped += 1
+            continue
+
+        split_name = _assign_split(file_time, train_range, val_range, test_range)
+        if split_name is None:
+            print(f"Skipping {filename}: no matching date range ({file_time.date()})")
+            skipped += 1
+            continue
+
+        src = os.path.join(input_dir, filename)
+        dst = os.path.join(output_dir, split_name, filename)
+        if os.path.exists(dst):
+            skipped += 1
+            continue
+
+        (shutil.copy2 if copy_files else shutil.move)(src, dst)
+        moved += 1
+
+    action = "copied" if copy_files else "moved"
+    print(f"Done: {moved} files {action}, {skipped} skipped")
+    for split_name in ("train", "val", "test"):
+        n = len(os.listdir(os.path.join(output_dir, split_name)))
+        print(f"  {split_name}: {n} files")

download/download_sdo.py

@@ -91,12 +91,20 @@ class SDODownloader:
         obs = pd.to_datetime(date_str, errors='coerce', utc=True).dt.tz_localize(None)
         header['_diff'] = (obs - date).abs()
 
+        # QUALITY comes back from JSOC as a numeric keyword, but drms occasionally
+        # yields it as object/string dtype — compare on the coerced numeric value
+        # so a dtype quirk can't silently make every frame look "not quality-0"
+        # (which would fall through to accepting whatever frame is closest,
+        # calibration/dark frames included).
+        quality_numeric = pd.to_numeric(header['QUALITY'], errors='coerce')
+
         rows = []
         for wl in wavelengths:
-            sub = header[header['WAVELNTH'] == int(wl)]
+            mask = header['WAVELNTH'] == int(wl)
+            sub = header[mask]
             if len(sub) == 0:
                 continue
-            good = sub[(sub['QUALITY'] == 0) | sub['QUALITY'].isna()]
+            good = sub[(quality_numeric[mask] == 0) | quality_numeric[mask].isna()]
             if len(good) > 0:
                 sub = good
             else:
@@ -173,13 +181,27 @@ class SDODownloader:
 
             header['DATE_OBS'] = header['DATE__OBS']
             header = header_to_fits(MetaDict(header))
-            with fits.open(tmp_path, 'update') as f:
-                hdr = f[1].header
-                for k, v in header.items():
-                    if pd.isna(v):
-                        continue
-                    hdr[k] = v
-                f.verify('silentfix')
+
+            # JSOC serves this segment as a tile-compressed CompImageHDU whose
+            # correct physical values depend on internal per-tile scaling
+            # (ZSCALE/ZZERO), not just the outer BSCALE/BZERO placeholder.
+            # Opening it in 'update' mode and writing header cards in place
+            # forces astropy to re-derive that scaling from the (unrelated)
+            # placeholder values, silently corrupting the pixel data — most
+            # visible on 94A, whose narrow true dynamic range gets clipped to
+            # a flat value after calibration. Instead: read the correctly
+            # decompressed data once, then write a fresh plain (uncompressed)
+            # FITS file with the merged header — never re-touch the original
+            # compressed HDU.
+            with fits.open(tmp_path) as f:
+                data = f[1].data
+                hdr = f[1].header.copy()
+            for k, v in header.items():
+                if pd.isna(v):
+                    continue
+                hdr[k] = v
+            fits.HDUList([fits.PrimaryHDU(), fits.ImageHDU(data=data, header=hdr)]).writeto(
+                tmp_path, overwrite=True)
 
             os.replace(tmp_path, map_path)
             return map_path

download/download_sxr.py

@@ -12,7 +12,7 @@ class SXRDownloader:
 
     def __init__(self, save_dir: str):
         self.save_dir = Path(save_dir)
-        self.save_dir.mkdir(exist_ok=True)
+        self.save_dir.mkdir(parents=True, exist_ok=True)
         self.used_g13_files = []
         self.used_g14_files = []
         self.used_g15_files = []

download/hf_download_config.yaml

@@ -46,6 +46,14 @@ subsample_frac: 0.1      # used only when subsample_n is null  (0.0–1.0)
 # Only used when subsample: true.
 shuffle_buffer_size: 10
 
+# -----------------------------------------------------------------------------
+# Already downloaded parquet files locally instead of streaming from HF Hub?
+# Point this at the root dir containing per-split subdirs (train/, validation/
+# or val/, test/) and streaming + HF Hub login are skipped entirely.
+# -----------------------------------------------------------------------------
+local_parquet_dir: null
+delete_parquet_after: false   # delete each parquet file once its rows are written
+
 # -----------------------------------------------------------------------------
 # Misc
 # -----------------------------------------------------------------------------

download/hugging_face_data_download.py

@@ -1,17 +1,18 @@
 """
-Download FOXES dataset from HuggingFace Hub
-============================================
-Reconstructs the local AIA / SXR directory layout expected by the pipeline:
+Get FOXES data from HuggingFace Hub: either stream it directly, or convert
+parquet files you've already downloaded locally. Reconstructs the local AIA /
+SXR directory layout expected by the pipeline:
     {aia_dir}/{split}/{filename}   — shape (7, 512, 512) float32 .npy
     {sxr_dir}/{split}/{filename}   — shape (N,) float32 .npy
 
-Uses streaming so parquet shards are fetched on-the-fly without loading the
-full dataset into RAM. A ThreadPoolExecutor overlaps disk writes with network
-fetches so the stream never stalls waiting for np.save to finish.
+Streaming mode fetches parquet shards on-the-fly without loading the full
+dataset into RAM; a ThreadPoolExecutor overlaps disk writes with network
+fetches so the stream never stalls waiting for np.save to finish. Local mode
+reads parquet files you already have on disk, using bulk/vectorized column
+conversion instead of a per-row Python loop.
 
 Usage:
-    python download/hugging_face_data_download.py \\
-        --config download/hf_download_config.yaml
+    python download/hugging_face_data_download.py --config download/hf_download_config.yaml
 """
 
 import argparse
@@ -19,8 +20,10 @@ import logging
 import os
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
 
 import numpy as np
+import pyarrow.parquet as pq
 import yaml
 from datasets import load_dataset
 from huggingface_hub import login
@@ -28,6 +31,7 @@ from huggingface_hub import login
 
 # HuggingFace uses "validation"; local pipeline directories use "val"
 HF_TO_LOCAL = {"validation": "val"}
+LOCAL_TO_HF = {v: k for k, v in HF_TO_LOCAL.items()}
 
 
 def load_config(path: str) -> dict:
@@ -35,6 +39,22 @@ def load_config(path: str) -> dict:
         return yaml.safe_load(f)
 
 
+def _write_arrays(filename: str, aia_arr: np.ndarray, sxr_arr: np.ndarray,
+                  aia_split_dir: str, sxr_split_dir: str) -> bool:
+    """Save pre-materialized arrays to disk. Returns True if written, False if already exists."""
+    aia_path = os.path.join(aia_split_dir, filename)
+    sxr_path = os.path.join(sxr_split_dir, filename)
+
+    if os.path.exists(aia_path) and os.path.exists(sxr_path):
+        return False
+
+    np.save(aia_path, aia_arr)
+    np.save(sxr_path, sxr_arr)
+    return True
+
+
+# --------------------------------------------------------------------------- streaming from HF Hub
+
 def resolve_n(repo_id: str, hf_split: str, cfg: dict) -> int | None:
     """
     Resolve how many rows to download. Returns an exact count, or None for all rows.
@@ -75,20 +95,6 @@ def build_dataset(repo_id: str, hf_split: str, cfg: dict, n: int | None):
     return ds
 
 
-def _write_arrays(filename: str, aia_arr: np.ndarray, sxr_arr: np.ndarray,
-                  aia_split_dir: str, sxr_split_dir: str) -> bool:
-    """Save pre-materialized arrays to disk. Returns True if written, False if already exists."""
-    aia_path = os.path.join(aia_split_dir, filename)
-    sxr_path = os.path.join(sxr_split_dir, filename)
-
-    if os.path.exists(aia_path) and os.path.exists(sxr_path):
-        return False
-
-    np.save(aia_path, aia_arr)
-    np.save(sxr_path, sxr_arr)
-    return True
-
-
 def download_split(hf_split: str, cfg: dict):
     local_split = HF_TO_LOCAL.get(hf_split, hf_split)
     aia_split_dir = os.path.join(cfg["aia_dir"], local_split)
@@ -156,20 +162,156 @@ def download_split(hf_split: str, cfg: dict):
     print(f"[{hf_split}] Done — {saved} saved, {skipped} skipped | {elapsed/60:.1f} min")
 
 
+# --------------------------------------------------------------------------- already-downloaded local parquet
+
+def _resolve_local_split_dir(local_parquet_dir: str, hf_split: str):
+    """Find the on-disk subdir for this split, trying both the HF name and its local alias."""
+    for candidate in (hf_split, HF_TO_LOCAL.get(hf_split), LOCAL_TO_HF.get(hf_split)):
+        if candidate is None:
+            continue
+        split_dir = os.path.join(local_parquet_dir, candidate)
+        if os.path.isdir(split_dir):
+            return split_dir
+    return None
+
+
+def _to_numpy_bulk(col, n):
+    """
+    Flatten a nested fixed-size list column directly to numpy, then reshape to
+    (n, *element_shape). Avoids per-row Python overhead (table.slice / .as_py()
+    on large nested arrays).
+    """
+    chunk = col.combine_chunks()
+    # Walk the array levels to recover per-element shape. FixedSizeList exposes
+    # list_size directly; regular List (which HF uses even for fixed-size
+    # arrays) stores a uniform stride in its offsets buffer — both give the
+    # same answer.
+    shape = []
+    vals = chunk
+    while hasattr(vals, "values"):
+        if hasattr(vals, "list_size"):   # FixedSizeList
+            shape.append(vals.list_size)
+        elif hasattr(vals, "offsets"):   # List — stride from offsets
+            shape.append(vals.offsets[1].as_py())
+        vals = vals.values
+    arr = vals.to_numpy(zero_copy_only=False).astype(np.float32)
+    return arr.reshape(n, *shape) if shape else arr.reshape(n, arr.size // n)
+
+
+def convert_local_parquet_split(parquet_dir: str, hf_split: str, cfg: dict):
+    local_split = HF_TO_LOCAL.get(hf_split, hf_split)
+    aia_split_dir = os.path.join(cfg["aia_dir"], local_split)
+    sxr_split_dir = os.path.join(cfg["sxr_dir"], local_split)
+    os.makedirs(aia_split_dir, exist_ok=True)
+    os.makedirs(sxr_split_dir, exist_ok=True)
+
+    parquet_files = sorted(Path(parquet_dir).glob("*.parquet"))
+    if not parquet_files:
+        print(f"[warn] No parquet files found in {parquet_dir}")
+        return
+
+    print(f"\n{'='*50}")
+    print(f"Converting split: {hf_split} -> local dir: {local_split}")
+    print(f"{'='*50}")
+    print(f"  Parquet dir: {parquet_dir} ({len(parquet_files)} files)")
+    print(f"  AIA -> {aia_split_dir}")
+    print(f"  SXR -> {sxr_split_dir}")
+
+    num_workers = cfg.get("num_workers", 8)
+    print_every = cfg.get("print_every", 500)
+    delete_after = cfg.get("delete_parquet_after", False)
+    saved = skipped = submitted = 0
+    start = time.time()
+
+    with ThreadPoolExecutor(max_workers=num_workers) as pool:
+        futures = {}
+
+        for pq_file in parquet_files:
+            table = pq.read_table(pq_file, columns=["filename", "aia_stack", "sxr_value"])
+            n_rows = len(table)
+            filenames = table.column("filename").to_pylist()
+
+            try:
+                aia_bulk = _to_numpy_bulk(table.column("aia_stack"), n_rows)
+                use_bulk_aia = True
+            except Exception:
+                aia_pylist = table.column("aia_stack").to_pylist()
+                use_bulk_aia = False
+
+            try:
+                sxr_bulk = _to_numpy_bulk(table.column("sxr_value"), n_rows)
+                use_bulk_sxr = True
+            except Exception:
+                sxr_pylist = table.column("sxr_value").to_pylist()
+                use_bulk_sxr = False
+
+            file_futures = []
+            for i, filename in enumerate(filenames):
+                aia_arr = aia_bulk[i] if use_bulk_aia else np.array(aia_pylist[i], dtype=np.float32)
+                sxr_arr = sxr_bulk[i] if use_bulk_sxr else np.array(sxr_pylist[i], dtype=np.float32)
+
+                fut = pool.submit(_write_arrays, filename, aia_arr, sxr_arr, aia_split_dir, sxr_split_dir)
+                file_futures.append(fut)
+                futures[fut] = submitted
+                submitted += 1
+
+                if submitted % print_every == 0:
+                    done = [f for f in futures if f.done()]
+                    for f in done:
+                        if f.result():
+                            saved += 1
+                        else:
+                            skipped += 1
+                        del futures[f]
+
+                    elapsed = time.time() - start
+                    rate = submitted / elapsed if elapsed > 0 else 0
+                    print(
+                        f"[{hf_split}] submitted={submitted} | saved={saved} skipped={skipped} | "
+                        f"{rate:.1f} rows/sec",
+                        flush=True,
+                    )
+
+            # Wait for all writes from this file to finish before deleting it
+            if delete_after:
+                for fut in as_completed(file_futures):
+                    pass
+                pq_file.unlink()
+                print(f"  Deleted {pq_file.name}", flush=True)
+
+        for fut in as_completed(futures):
+            if fut.result():
+                saved += 1
+            else:
+                skipped += 1
+
+    elapsed = time.time() - start
+    print(f"[{hf_split}] Done — {saved} saved, {skipped} skipped | {elapsed/60:.1f} min")
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Download FOXES data from HuggingFace Hub")
+    parser = argparse.ArgumentParser(description="Get FOXES data from HuggingFace Hub (stream or local parquet)")
     parser.add_argument("--config", default="download/hf_download_config.yaml")
     args = parser.parse_args()
-
     cfg = load_config(args.config)
 
-    login()
-
     splits = cfg.get("splits", ["train", "validation", "test"])
-    for split in splits:
-        download_split(split, cfg)
-
-    print("\nAll splits downloaded successfully.")
+    local_parquet_dir = cfg.get("local_parquet_dir")
+
+    if local_parquet_dir:
+        print(f"Converting local parquet files from {local_parquet_dir} (no HF Hub access needed)")
+        for split in splits:
+            split_dir = _resolve_local_split_dir(local_parquet_dir, split)
+            if split_dir is None:
+                print(f"[warn] Split dir not found under {local_parquet_dir}, skipping: {split}")
+                continue
+            convert_local_parquet_split(split_dir, split, cfg)
+    else:
+        login()
+        for split in splits:
+            download_split(split, cfg)
+
+    print("\nAll splits done.")
     # Silence HF/fsspec connection pool teardown — it throws [Errno 9] Bad file
     # descriptor during GC which is harmless but noisy.
     for name in ("huggingface_hub", "urllib3", "fsspec", "datasets"):

download/sdo_download_config.yaml

@@ -10,11 +10,11 @@
 #   python download/download_sdo.py --config download/sdo_download_config.yaml
 # =============================================================================
 
-download_dir: "/path/to/AIA_raw"
-email: "your.email@example.com"   # must be registered with JSOC: http://jsoc.stanford.edu/ajax/register_email.html
+download_dir: "/Volumes/T9/testing_foxes/AIA_raw"
+email: "ggoodwin5@gsu.edu"   # must be registered with JSOC: http://jsoc.stanford.edu/ajax/register_email.html
 
 start_date: "2023-08-01 00:00:00"
-end_date:   "2023-08-02 00:00:00"
-cadence: 60    # minutes between samples
+end_date:   "2023-12-01 00:00:00"
+cadence: 7200    # minutes between samples
 
 wavelengths: [94, 131, 171, 193, 211, 304, 335]

download/sxr_download_config.yaml

@@ -1,19 +1,19 @@
 # =============================================================================
 # Raw GOES/SXR Download Configuration
 # =============================================================================
-# Used by sxr_downloader.py to pull raw GOES XRS netCDF files via SunPy Fido.
+# Used by download_sxr.py to pull raw GOES XRS netCDF files via SunPy Fido.
 # This is the advanced/raw path — most users should use
 # hugging_face_data_download.py instead (see the README). Data downloaded
 # here needs the data/ scripts to become the paired .npy format the model
 # expects.
 #
 # Usage:
-#   python download/sxr_downloader.py --config download/sxr_download_config.yaml
+#   python download/download_sxr.py --config download/sxr_download_config.yaml
 # =============================================================================
 
-save_dir: "/path/to/GOES_raw"
+save_dir: "/Volumes/T9/testing_foxes/GOES_raw"
 
 start_date: "2023-08-01"
-end_date:   "2023-08-02"
+end_date:   "2023-08-31"
 
 max_workers: 4

forecasting/dataset.py

@@ -4,6 +4,8 @@ import pandas as pd
 import torch
 import numpy as np
 from pathlib import Path
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader
 import glob
 
 
@@ -170,3 +172,56 @@ class AIAGOESDataset(torch.utils.data.Dataset):
         return timestamp
 
 
+class AIAGOESDataModule(LightningDataModule):
+    """
+    PyTorch Lightning DataModule wiring up train/val/test AIAGOESDataset splits.
+    Used by train.py.
+
+    Parameters
+    ----------
+    aia_train_dir, aia_val_dir, aia_test_dir : str
+        Directories of AIA .npy files for each split.
+    sxr_train_dir, sxr_val_dir, sxr_test_dir : str
+        Directories of SXR .npy files for each split.
+    sxr_norm : np.ndarray
+        (mean, std) used to log-normalize SXR targets.
+    batch_size, num_workers : int
+    wavelengths : list of int
+    """
+
+    def __init__(self, aia_train_dir, aia_val_dir, aia_test_dir, sxr_train_dir, sxr_val_dir, sxr_test_dir,
+                 sxr_norm, batch_size=64, num_workers=4, wavelengths=[94, 131, 171, 193, 211, 304, 335]):
+        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.batch_size = batch_size
+        self.num_workers = num_workers
+        self.wavelengths = wavelengths
+
+    def setup(self, stage=None):
+        transform = SXRLogNormTransform(self.sxr_norm[0], self.sxr_norm[1])
+        self.train_ds = AIAGOESDataset(aia_dir=self.aia_train_dir, sxr_dir=self.sxr_train_dir,
+                                       sxr_transform=transform, wavelengths=self.wavelengths)
+        self.val_ds = AIAGOESDataset(aia_dir=self.aia_val_dir, sxr_dir=self.sxr_val_dir,
+                                     sxr_transform=transform, wavelengths=self.wavelengths)
+        self.test_ds = AIAGOESDataset(aia_dir=self.aia_test_dir, sxr_dir=self.sxr_test_dir,
+                                      sxr_transform=transform, wavelengths=self.wavelengths)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True,
+                          num_workers=self.num_workers, prefetch_factor=4 if self.num_workers else None)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_ds, batch_size=self.batch_size, shuffle=False,
+                          num_workers=self.num_workers, prefetch_factor=4 if self.num_workers else None)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_ds, batch_size=self.batch_size, shuffle=False,
+                          num_workers=self.num_workers, prefetch_factor=1 if self.num_workers else None)
+
+

forecasting/evaluation.py

@@ -474,7 +474,7 @@ def load_evaluation_config(config_path):
         config_data = yaml.load(stream, Loader=yaml.SafeLoader)
 
     # Resolve variable substitutions
-    config_data = resolve_config_variables(config_data)
+    config_data: dict = resolve_config_variables(config_data)
     return config_data
 
 
@@ -488,7 +488,7 @@ def main():
     import argparse
 
     parser = argparse.ArgumentParser(description='Run FOXES solar flare evaluation')
-    parser.add_argument('-config', type=str, default='evaluation_config.yaml',
+    parser.add_argument('--config', type=str, default='evaluation_config.yaml',
                         help='Path to evaluation config YAML file')
     args = parser.parse_args()
 

forecasting/evaluation_config.yaml

@@ -3,7 +3,7 @@
 # =============================================================================
 # Used by evaluation.py to compute metrics and generate plots.
 #
-# Usage: python evaluation.py -config evaluation_config.yaml
+# Usage: python evaluation.py --config evaluation_config.yaml
 # =============================================================================
 
 model_predictions:

forecasting/inference.py

@@ -143,8 +143,6 @@ def evaluate_model_on_dataset(model, dataset, batch_size=16, times=None, config_
                 flux_contributions = None
 
             current_batch_size = predictions.shape[0]
-            batch_weights = []
-            batch_flux_contributions = []
             
             # Process each sample in the batch to reduce memory footprint
             for i in range(current_batch_size):
@@ -319,14 +317,14 @@ def main():
         return recursive_substitute(config_dict, variables)
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('-config', type=str, default='inference_config.yaml', required=True,
+    parser.add_argument('--config', type=str, default='inference_config.yaml', required=True,
                         help='Path to the inference configuration YAML file.')
     args = parser.parse_args()
 
     with open(args.config, 'r') as stream:
         config_data = yaml.load(stream, Loader=yaml.SafeLoader)
 
-    config_data = resolve_config_variables(config_data)
+    config_data: dict = resolve_config_variables(config_data)
 
     model_params = config_data.get('model_params', {})
     input_size = model_params.get('input_size', 512)

forecasting/inference_config.yaml

@@ -5,7 +5,7 @@
 # AIA data and write out SXR predictions (+ optional attention/flux maps).
 #
 # Usage:
-#   python forecasting/inference.py -config forecasting/inference_config.yaml
+#   python forecasting/inference.py --config forecasting/inference_config.yaml
 # =============================================================================
 
 wavelengths: [94, 131, 171, 193, 211, 304, 335]

forecasting/trained_weights_and_normalization/localized-attention-mask.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df37bef4ac95a3c14e777221ce1fd3fbdf56da78a192d7418303848d84aa2f46
+size 224141974

forecasting/trained_weights_and_normalization/no-attention-mask.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d6965e5b7c1f41565b1eebf059ce5506d6e6c8b8483f7b569a854f1902ae45b
+size 89939767

training/callbacks.py

@@ -0,0 +1,226 @@
+"""
+PyTorch Lightning callbacks for visualizing training progress: predicted vs.
+true SXR flux, and Vision Transformer attention maps, logged to Weights & Biases.
+
+Used by train.py — not meant to be run standalone.
+"""
+import wandb
+from pytorch_lightning import Callback
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from forecasting.model import unnormalize_sxr
+
+
+class ImagePredictionLogger_SXR(Callback):
+    """
+    PyTorch Lightning callback for logging AIA input images and corresponding
+    true vs predicted Soft X-Ray (SXR) flux values to Weights & Biases (wandb).
+
+    This helps monitor model performance across validation epochs by
+    comparing predicted vs. ground-truth flare intensities.
+    """
+
+    def __init__(self, data_samples, sxr_norm):
+        """
+        Initialize callback with validation samples and normalization parameters.
+
+        Parameters
+        ----------
+        data_samples : list
+            List of validation samples (AIA image, SXR target pairs).
+        sxr_norm : np.ndarray
+            Normalization statistics used to unnormalize predicted flux values.
+        """
+        super().__init__()
+        self.data_samples = data_samples
+        self.sxr_norm = sxr_norm
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        """
+        Log scatter plots comparing predicted and true SXR flux values
+        at the end of each validation epoch.
+        """
+        true_sxr = []
+        pred_sxr = []
+
+        for aia, target in self.data_samples:
+            aia = aia.to(pl_module.device).unsqueeze(0)
+            # forward() always returns a tuple (global_flux_raw, ...); we only need the flux.
+            pred, *_ = pl_module(aia, return_attention=False)
+            pred_sxr.append(pred.item())
+            true_sxr.append(target.item())
+
+        true_unorm = unnormalize_sxr(np.array(true_sxr, dtype=np.float32), self.sxr_norm)
+        pred_unnorm = unnormalize_sxr(np.array(pred_sxr, dtype=np.float32), self.sxr_norm)
+
+        fig1 = self.plot_aia_sxr(true_unorm, pred_unnorm)
+        trainer.logger.experiment.log({"Soft X-ray flux plots": wandb.Image(fig1)})
+        plt.close(fig1)
+
+    # Flare-class range: A-class starts at 1e-8 W/m^2, X10 is 10x the 1e-4 X-class threshold.
+    AXIS_MIN = 1e-8
+    AXIS_MAX = 1e-3
+
+    def plot_aia_sxr(self, val_sxr, pred_sxr):
+        """Log-log parity plot: predicted vs. true SXR flux, with a 1:1 reference line."""
+        fig, ax = plt.subplots(1, 1, figsize=(4, 4))
+
+        ax.plot([self.AXIS_MIN, self.AXIS_MAX], [self.AXIS_MIN, self.AXIS_MAX],
+                color='gray', linestyle='--', linewidth=1, label='Perfect prediction')
+        ax.scatter(val_sxr, pred_sxr, color='blue', alpha=0.7,s=10, label='Predictions')
+
+        ax.set_xscale('log')
+        ax.set_yscale('log')
+        ax.set_xlim(self.AXIS_MIN, self.AXIS_MAX)
+        ax.set_ylim(self.AXIS_MIN, self.AXIS_MAX)
+        ax.set_xlabel("True SXR flux [W/m$^2$]")
+        ax.set_ylabel("Predicted SXR flux [W/m$^2$]")
+        ax.legend()
+        fig.tight_layout()
+        return fig
+
+
+class AttentionMapCallback(Callback):
+    """
+    PyTorch Lightning callback for visualizing transformer attention maps
+    during validation epochs.
+
+    Supports CLS-token-based and local patch attention visualization.
+    """
+
+    def __init__(self, log_every_n_epochs=1, num_samples=4, patch_size=8, use_local_attention=False):
+        """
+        Initialize callback.
+
+        Parameters
+        ----------
+        log_every_n_epochs : int
+            Frequency of logging attention maps.
+        num_samples : int
+            Number of samples to visualize per epoch.
+        patch_size : int
+            Patch size used in the Vision Transformer.
+        use_local_attention : bool
+            If True, visualize local attention patterns instead of CLS attention.
+        """
+        super().__init__()
+        self.patch_size = patch_size
+        self.log_every_n_epochs = log_every_n_epochs
+        self.num_samples = num_samples
+        self.use_local_attention = use_local_attention
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        """Trigger visualization of attention maps at the end of validation epochs."""
+        if trainer.current_epoch % self.log_every_n_epochs == 0:
+            self._visualize_attention(trainer, pl_module)
+
+    def _visualize_attention(self, trainer, pl_module):
+        """Generate and log attention maps from the model's attention weights."""
+        val_dataloader = trainer.val_dataloaders
+        if val_dataloader is None:
+            return
+
+        pl_module.eval()
+        with torch.no_grad():
+            batch = next(iter(val_dataloader))
+            imgs, labels = batch
+            imgs = imgs[:self.num_samples].to(pl_module.device)
+
+            # forward(return_attention=True) -> (global_flux_raw, attention_weights, patch_flux_raw)
+            _, attention_weights, patch_flux_raw = pl_module(imgs, return_attention=True)
+
+            for sample_idx in range(min(self.num_samples, imgs.size(0))):
+                fig = self._plot_attention_map(
+                    imgs[sample_idx],
+                    attention_weights,
+                    sample_idx,
+                    trainer.current_epoch,
+                    patch_size=self.patch_size,
+                    patch_flux=patch_flux_raw[sample_idx] if patch_flux_raw is not None else None
+                )
+                trainer.logger.experiment.log({"Attention plots": wandb.Image(fig)})
+                plt.close(fig)
+
+    def _plot_attention_map(self, image, attention_weights, sample_idx, epoch, patch_size, patch_flux=None):
+        """Plot and return a visualization of the attention heatmaps for a single image."""
+        img_np = image.cpu().numpy()
+        if len(img_np.shape) == 3 and img_np.shape[0] in [1, 3]:
+            img_np = np.transpose(img_np, (1, 2, 0))
+
+        H, W = img_np.shape[:2]
+        grid_h, grid_w = H // patch_size, W // patch_size
+
+        last_layer_attention = attention_weights[-1]
+        sample_attention = last_layer_attention[sample_idx]
+        avg_attention = sample_attention.mean(dim=0)
+
+        if self.use_local_attention:
+            # Spatial center of the grid, not the middle of the flattened sequence
+            # (those only coincide when grid_w == 1) — row-major flatten: idx = row*grid_w + col.
+            center_patch_idx = (grid_h // 2) * grid_w + (grid_w // 2)
+            center_attention = avg_attention[center_patch_idx, :].cpu()
+            avg_attention_map = avg_attention.mean(dim=0).cpu()
+            attention_map = avg_attention_map.reshape(grid_h, grid_w)
+            center_map = center_attention.reshape(grid_h, grid_w)
+        else:
+            cls_attention = avg_attention[0, 1:].cpu()
+            attention_map = cls_attention.reshape(grid_h, grid_w)
+            center_map = None
+
+        if len(img_np[0, 0, :]) >= 6:
+            rgb_channels = [0, 2, 4]
+            img_display = np.stack([(img_np[:, :, i] + 1) / 2 for i in rgb_channels], axis=2)
+            img_display = np.clip(img_display, 0, 1)
+        else:
+            img_display = (img_np[:, :, 0] + 1) / 2
+            img_display = np.stack([img_display] * 3, axis=2)
+
+        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+        fig.suptitle(f'Attention Visualization - Epoch {epoch}, Sample {sample_idx}', fontsize=16)
+
+        axes[0, 0].imshow(img_display)
+        axes[0, 0].set_title('Original Image')
+        axes[0, 0].axis('off')
+
+        im1 = axes[0, 1].imshow(attention_map, cmap='hot', interpolation='nearest')
+        axes[0, 1].set_title('Attention Map')
+        axes[0, 1].axis('off')
+        plt.colorbar(im1, ax=axes[0, 1])
+
+        axes[0, 2].imshow(img_display)
+        axes[0, 2].imshow(attention_map, cmap='hot', alpha=0.6, interpolation='nearest')
+        axes[0, 2].set_title('Attention Overlay')
+        axes[0, 2].axis('off')
+
+        if center_map is not None:
+            im2 = axes[1, 0].imshow(center_map, cmap='hot', interpolation='nearest')
+            axes[1, 0].set_title('Center Patch Attention')
+            axes[1, 0].axis('off')
+            plt.colorbar(im2, ax=axes[1, 0])
+        else:
+            axes[1, 0].text(0.5, 0.5, 'Center attention\nnot available',
+                           ha='center', va='center', transform=axes[1, 0].transAxes)
+            axes[1, 0].set_title('Center Patch Attention')
+            axes[1, 0].axis('off')
+
+        if patch_flux is not None:
+            patch_flux_np = patch_flux.cpu().numpy().reshape(grid_h, grid_w)
+            im3 = axes[1, 1].imshow(patch_flux_np, cmap='viridis', interpolation='nearest')
+            axes[1, 1].set_title('Patch Flux')
+            axes[1, 1].axis('off')
+            plt.colorbar(im3, ax=axes[1, 1])
+        else:
+            axes[1, 1].text(0.5, 0.5, 'Patch flux\nnot available',
+                           ha='center', va='center', transform=axes[1, 1].transAxes)
+            axes[1, 1].set_title('Patch Flux')
+            axes[1, 1].axis('off')
+
+        axes[1, 2].hist(attention_map.flatten(), bins=50, alpha=0.7)
+        axes[1, 2].set_title('Attention Distribution')
+        axes[1, 2].set_xlabel('Attention Weight')
+        axes[1, 2].set_ylabel('Frequency')
+
+        plt.tight_layout()
+        return fig

training/train.py

@@ -0,0 +1,224 @@
+"""
+Training script for the AIA-GOES multimodal solar flare forecasting model using PyTorch Lightning.
+
+This script:
+1. Loads configuration from a YAML file with variable substitution (e.g., ${base_dir} references).
+2. Initializes the AIA-GOES DataModule.
+3. Configures logging with Weights & Biases.
+4. Builds and trains a Vision Transformer (ViTLocal) model.
+5. Optionally computes dynamic base class weights for flare categories (Quiet, C, M, X).
+6. Saves model checkpoints (.ckpt).
+
+Usage:
+    python training/train.py --config training/train_config.yaml
+"""
+
+import argparse
+import os
+import re
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import wandb
+import yaml
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.loggers import WandbLogger
+
+PROJECT_ROOT = Path(__file__).parent.parent.absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from training.callbacks import AttentionMapCallback, ImagePredictionLogger_SXR
+from forecasting.dataset import AIAGOESDataModule
+from forecasting.model import ViTLocal, unnormalize_sxr
+
+
+def resolve_config_variables(config_dict):
+    """
+    Recursively resolve ${variable} references within the config.
+
+    This function processes configuration dictionaries to substitute variable
+    references of the form ${variable_name} with their actual values defined
+    elsewhere in the configuration.
+    """
+    variables = {}
+    for key, value in config_dict.items():
+        if isinstance(value, str) and not value.startswith('${'):
+            variables[key] = value
+
+    def substitute_value(value, variables):
+        if isinstance(value, str):
+            pattern = r'\$\{([^}]+)\}'
+            for match in re.finditer(pattern, value):
+                var_name = match.group(1)
+                if var_name in variables:
+                    value = value.replace(f'${{{var_name}}}', variables[var_name])
+        return value
+
+    def recursive_substitute(obj, variables):
+        if isinstance(obj, dict):
+            return {k: recursive_substitute(v, variables) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [recursive_substitute(item, variables) for item in obj]
+        else:
+            return substitute_value(obj, variables)
+
+    return recursive_substitute(config_dict, variables)
+
+
+def get_base_weights(data_module, sxr_norm):
+    """
+    Compute inverse-frequency weights for flare classes based on training data.
+
+    The weights help balance loss contributions from imbalanced flare categories
+    by making rare classes (M/X flares) count for more in the loss.
+
+    Parameters
+    ----------
+    data_module : AIAGOESDataModule
+        Initialized DataModule providing the train_dataloader.
+    sxr_norm : np.ndarray
+        Normalization parameters for SXR.
+
+    Returns
+    -------
+    dict
+        Class weights for quiet, C, M, and X classes.
+    """
+    print("Calculating base weights from training data...")
+    c_threshold, m_threshold, x_threshold = 1e-6, 1e-5, 1e-4
+
+    quiet_count = c_count = m_count = x_count = total = 0
+    train_loader = data_module.train_dataloader()
+    print(f"Processing {len(train_loader)} batches...")
+
+    for batch_idx, (aia_batch, sxr_batch) in enumerate(train_loader):
+        if batch_idx % 50 == 0:
+            print(f"Processed {batch_idx}/{len(train_loader)} batches...")
+
+        sxr_un = unnormalize_sxr(sxr_batch, sxr_norm)
+        sxr_un_flat = sxr_un.reshape(-1)
+
+        total += len(sxr_un_flat)
+        quiet_count += (sxr_un_flat < c_threshold).sum()
+        c_count += ((sxr_un_flat >= c_threshold) & (sxr_un_flat < m_threshold)).sum()
+        m_count += ((sxr_un_flat >= m_threshold) & (sxr_un_flat < x_threshold)).sum()
+        x_count += (sxr_un_flat >= x_threshold).sum()
+
+    quiet_count, c_count, m_count, x_count = (max(c, 1) for c in (quiet_count, c_count, m_count, x_count))
+
+    weights = {
+        'quiet': total / quiet_count,
+        'c_class': total / c_count,
+        'm_class': total / m_count,
+        'x_class': total / x_count,
+    }
+    print(f"Total samples: {total}")
+    print(f"Quiet: {quiet_count} (weight {weights['quiet']:.4f}), "
+          f"C: {c_count} (weight {weights['c_class']:.4f}), "
+          f"M: {m_count} (weight {weights['m_class']:.4f}), "
+          f"X: {x_count} (weight {weights['x_class']:.4f})")
+    return weights
+
+
+def resolve_devices(gpu_config):
+    """Resolve accelerator/devices/strategy from the gpu_ids config value."""
+    if gpu_config == -1:
+        print("Using CPU for training")
+        return "cpu", 1, "auto"
+
+    if not torch.cuda.is_available():
+        print("No GPUs available, falling back to CPU")
+        return "cpu", 1, "auto"
+
+    if gpu_config == "all":
+        print(f"Using all available GPUs ({torch.cuda.device_count()} GPUs)")
+        return "gpu", -1, "auto"
+    if isinstance(gpu_config, list):
+        print(f"Using GPUs: {gpu_config}")
+        return "gpu", gpu_config, "auto"
+    print(f"Using GPU {gpu_config}")
+    return "gpu", [gpu_config], "auto"
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Train the FOXES ViTLocal model.')
+    parser.add_argument('--config', type=str, default='training/train_config.yaml', required=True,
+                        help='Path to train_config.yaml')
+    args = parser.parse_args()
+
+    with open(args.config, 'r') as stream:
+        config_data = yaml.load(stream, Loader=yaml.SafeLoader)
+    config_data: dict = resolve_config_variables(config_data)
+
+    print("Resolved paths:")
+    print(f"AIA dir: {config_data['data']['aia_dir']}")
+    print(f"SXR dir: {config_data['data']['sxr_dir']}")
+    print(f"Checkpoints dir: {config_data['data']['checkpoints_dir']}")
+
+    sxr_norm = np.load(config_data['data']['sxr_norm_path'])
+    wavelengths = config_data['wavelengths']
+
+    data_module = AIAGOESDataModule(
+        aia_train_dir=config_data['data']['aia_dir'] + "/train",
+        aia_val_dir=config_data['data']['aia_dir'] + "/val",
+        aia_test_dir=config_data['data']['aia_dir'] + "/test",
+        sxr_train_dir=config_data['data']['sxr_dir'] + "/train",
+        sxr_val_dir=config_data['data']['sxr_dir'] + "/val",
+        sxr_test_dir=config_data['data']['sxr_dir'] + "/test",
+        batch_size=config_data['batch_size'],
+        num_workers=min(8, os.cpu_count()),
+        sxr_norm=sxr_norm,
+        wavelengths=wavelengths,
+    )
+    data_module.setup()
+
+    wandb_logger = WandbLogger(
+        entity=config_data['wandb']['entity'],
+        project=config_data['wandb']['project'],
+        job_type=config_data['wandb']['job_type'],
+        tags=config_data['wandb']['tags'],
+        name=config_data['wandb']['run_name'],
+        notes=config_data['wandb']['notes'],
+        config=config_data,
+    )
+
+    # Callbacks
+    total_n_valid = len(data_module.val_ds)
+    plot_samples = [data_module.val_ds[i] for i in range(0, total_n_valid, max(1, total_n_valid // 4))]
+    sxr_plot_callback = ImagePredictionLogger_SXR(plot_samples, sxr_norm)
+    patch_size = config_data.get('vit_architecture', {}).get('patch_size', 16)
+    attention_callback = AttentionMapCallback(patch_size=patch_size, use_local_attention=True)
+
+    base_weights = get_base_weights(data_module, sxr_norm) if config_data.get('calculate_base_weights') else None
+    model = ViTLocal(model_kwargs=config_data['vit_architecture'], sxr_norm=sxr_norm, base_weights=base_weights)
+
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=config_data['data']['checkpoints_dir'],
+        monitor='val_total_loss',
+        mode='min',
+        save_top_k=10,
+        filename=f"{config_data['wandb']['run_name']}-{{epoch:02d}}-{{val_total_loss:.4f}}",
+    )
+
+    accelerator, devices, strategy = resolve_devices(config_data.get('gpu_ids', -1))
+
+    trainer = Trainer(
+        default_root_dir=config_data['data']['checkpoints_dir'],
+        accelerator=accelerator,
+        devices=devices,
+        strategy=strategy,
+        max_epochs=config_data['epochs'],
+        callbacks=[sxr_plot_callback, attention_callback, checkpoint_callback],
+        logger=wandb_logger,
+        log_every_n_steps=10,
+    )
+    trainer.fit(model, data_module)
+    wandb.finish()
+    print(f"Training complete. Checkpoints saved to {config_data['data']['checkpoints_dir']}")
+
+
+if __name__ == '__main__':
+    main()

training/train_config.yaml

@@ -0,0 +1,61 @@
+# =============================================================================
+# FOXES Training Configuration
+# =============================================================================
+# Used by train.py to train the ViTLocal model on paired AIA/SXR .npy data.
+# Expects aia_dir/sxr_dir to each have train/val/test subfolders (that's what
+# data/build_dataset.py produces).
+#
+# Usage:
+#   python training/train.py --config training/train_config.yaml
+# =============================================================================
+wavelengths: [94, 131, 171, 193, 211, 304, 335]
+
+# GPU configuration:
+#   0 (or any int)  - that single GPU
+#   [0, 1]          - those specific GPUs
+#   "all"           - every available GPU
+#   -1              - CPU only
+gpu_ids: -1
+
+batch_size: 6
+epochs: 150
+
+# Compute inverse-frequency loss weights (quiet/C/M/X) from the training data
+# before starting. Takes one extra pass over the training set.
+calculate_base_weights: false
+
+vit_architecture:
+  embed_dim: 256
+  num_channels: 7
+  num_classes: 1
+  patch_size: 8
+  num_patches: 4096
+  hidden_dim: 1024
+  num_heads: 8
+  num_layers: 8
+  dropout: 0.1
+  learning_rate: 0.0001
+  # Self-attention masking - one of:
+  #   inverted - the released model's behavior: each patch attends to DISTANT
+  #              patches (the flipped local-attention mask it was trained with)
+  #   local    - true local attention: each patch attends to its neighborhood
+  #   none     - standard full/global attention (no masking)
+  mask_mode: inverted
+  local_window: 9   # neighbourhood side length (patches), used by inverted/local
+
+data:
+  aia_dir: "/Volumes/T9/testing_foxes/AIA_processed"
+  sxr_dir: "/Volumes/T9/testing_foxes/SXR_processed"
+  sxr_norm_path: "/Volumes/T9/testing_foxes/normalized_sxr.npy"
+  checkpoints_dir: "/Volumes/T9/chkpt/"  # update to actual checkpoint
+
+wandb:
+  entity: "ggoodwin5-georgia-state-university"  # your W&B username or team name
+  project: FOXES
+  job_type: training
+  tags:
+    - aia
+    - sxr
+    - regression
+  run_name: run_1
+  notes: AIA to SXR translation