sdofmv2 0.1.0__tar.gz

sdofmv2-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Joseph Gallego
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sdofmv2-0.1.0/PKG-INFO

@@ -0,0 +1,149 @@
+Metadata-Version: 2.4
+Name: sdofmv2
+Version: 0.1.0
+Summary: Solar phenomena prediction models
+Author-email: Joseph Gallego <joaggi@gmail.com>, Daniela Martin <dmartinvega@gmail.com>, Jinsu Hong <jinsuhong.knight@gmail.com>
+License: MIT
+Project-URL: Repository, https://github.com/Joaggi/sdofmv2
+Project-URL: Issues, https://github.com/Joaggi/sdofmv2/issues
+Keywords: foundation model,solar physics,deep learning,space weather
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.5.1
+Requires-Dist: lightning>=2.6.0
+Requires-Dist: numpy>=2.3.5
+Requires-Dist: pandas>=2.3.3
+Requires-Dist: transformers>=4.57.3
+Requires-Dist: sunpy>=7.0.4
+Requires-Dist: astropy>=6.0
+Requires-Dist: timm>=1.0.22
+Requires-Dist: einops>=0.8.1
+Requires-Dist: hydra-core>=1.3.2
+Requires-Dist: wandb>=0.23.1
+Dynamic: license-file
+
+# SDO FM v2: [Full Title of the Project/Model]
+
+[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-%23EE4C2C.svg?style=flat&logo=PyTorch&logoColor=white)](https://pytorch.org/)
+[![PyTorch Lightning](https://img.shields.io/badge/PyTorch_Lightning-%23792EE5.svg?style=flat&logo=pytorchlightning&logoColor=white)](https://lightning.ai/docs/pytorch/stable/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+## Introduction
+**SDOFMv2** is an advanced multi-instrument foundation model designed to analyze Solar Dynamics Observatory (SDO) data and drive large-scale, data-driven heliophysics research. Building upon the original SDOFM framework, this version addresses previous limitations like restricted temporal coverage and reconstruction artifacts to significantly improve spatial coherence and global consistency.
+
+![Model architecture](sdofmv2.svg) 
+*A Masked Autoencoder (MAE) based on a Vision Transformer (ViT) architecture is utilized for pretraining. During this phase, a% of the image patches are masked, while the remaining (100 - a)% are processed by the encoder. The decoder block then reconstructs all patches, optimized via a customized loss function.*
+
+---
+
+## Getting Started
+
+### Prerequisites
+* Linux or macOS
+* Python 3.11+
+* NVIDIA GPU + CUDA toolkit (Recommended for training)
+
+### Environment Setup
+We recommend using `mamba` to manage dependencies. 
+
+> **Important Hardware Note:** > The `sdofmv2_environment.yml` file is configured for **CUDA 12.8** by default. If your hardware or drivers require a different CUDA version (e.g., CUDA 11.8), please open `sdofmv2_environment.yml` and modify the `pip` section at the bottom to match your system (e.g., change `cu128` to `cu118`) before running the setup commands.
+
+**Using Mamba:**
+```bash
+# Clone the repository
+git clone [https://github.com/Joaggi/sdofmv2.git](https://github.com/Joaggi/sdofmv2.git)
+cd sdofmv2
+
+# Create and activate the environment (This automatically installs PyTorch and the local package)
+mamba env create -f sdofmv2_environment.yml
+mamba activate sdofmv2
+```
+---
+
+## Repository Structure
+
+```text
+.
+├── configs/                # YAML configurations for experiments
+│   ├── downstream/         # Configs for downstream tasks (F10.7, solar wind)
+│   └── pretrain/           # Configs for MAE pretraining (AIA, HMI)
+├── notebooks/              # Jupyter notebooks for analysis and visualization
+│   ├── analysis/           # Attention maps, PCA, and masking analysis
+│   └── downstream_apps/    # How to use downstream scripts (Notebooks) for F10.7 and missing data applications
+├── scripts/                # Executable scripts for training and testing
+│   ├── pretrain.py         # Main pretraining script
+│   ├── finetuning_*.py     # Scripts for downstream finetuning
+│   └── test.py             # Script for evaluating checkpoints
+├── src/                    # Core source code package
+│   └── sdofmv2/
+│       ├── core/           # Base model architectures and modules
+│       ├── tasks/          # PyTorch Lightning modules (model & data module) for downstream tasks
+│       └── utils/          # Helper functions, physical constants and metrics
+├── pyproject.toml          # Project metadata and build dependencies
+└── sdofmv2_environment.yml # Mamba environment definition file
+```
+
+---
+
+## How to Use
+
+*(Note: It is recommended to run all scripts from the root directory of the repository so that file paths to `configs/` and `src/` resolve correctly.)*
+
+### 1. Data Preparation
+Before training or running inference, you need to prepare the dataset. 
+[Explain where to download the data, or provide a command if you have a script for it.]
+```bash
+python scripts/download_data_cache.py --target_dir ./assets/
+```
+
+### 2. Training the Model
+To train the model from scratch, execute the pretraining script and pass the relevant configuration file.
+```bash
+python scripts/pretrain.py --config-name pretrain_mae_AIA.yaml
+```
+
+### 3. Inference and Evaluation
+To evaluate a pre-trained checkpoint on the test set:
+```bash
+python scripts/test.py --config-name pretrain_mae_AIA.yaml
+```
+
+### 4. Downstream Finetuning
+To finetune the model on a specific downstream task (e.g., solar wind forecasting):
+```bash
+python scripts/finetuning_solarwind.py --config-name finetune_solarwind_config.yaml
+```
+---
+
+## Results & Visualizations
+[Include a brief summary of the model's performance. You can add a table of metrics or a sample plot showing predictions vs. ground truth.]
+
+![Sample Visualization](notebooks/analysis/SDOFMv2_AIA_results_exp.png) 
+*The first row displays the original ground-truth images. The second and third rows show the model's reconstructed images using masking ratios of 0% and 50%, respectively.*
+
+---
+
+## Citation
+If you find this repository or model useful in your academic research, please consider citing our work:
+
+```bibtex
+@misc{sdofmv2,
+  author = {Hong, Jinsu and Martin, Daniela and Gallego, Joseph},
+  title = {SDOFMv2: A Multi-Instrument Foundation Model for the Solar Dynamics Observatory with Transferable Downstream Applications},
+  year = {2026},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{[https://github.com/Joaggi/sdofmv2](https://github.com/Joaggi/sdofmv2)}},
+  note = {Jinsu Hong, Daniela Martin, and Joseph Gallego contributed equally to this work}
+}
+```
+
+## Contributing
+Contributions, bug reports, and feature requests are welcome! Please feel free to check the [issues page](https://github.com/Joaggi/sdofmv2/issues) or submit a pull request.

sdofmv2-0.1.0/README.md

@@ -0,0 +1,119 @@
+# SDO FM v2: [Full Title of the Project/Model]
+
+[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-%23EE4C2C.svg?style=flat&logo=PyTorch&logoColor=white)](https://pytorch.org/)
+[![PyTorch Lightning](https://img.shields.io/badge/PyTorch_Lightning-%23792EE5.svg?style=flat&logo=pytorchlightning&logoColor=white)](https://lightning.ai/docs/pytorch/stable/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+## Introduction
+**SDOFMv2** is an advanced multi-instrument foundation model designed to analyze Solar Dynamics Observatory (SDO) data and drive large-scale, data-driven heliophysics research. Building upon the original SDOFM framework, this version addresses previous limitations like restricted temporal coverage and reconstruction artifacts to significantly improve spatial coherence and global consistency.
+
+![Model architecture](sdofmv2.svg) 
+*A Masked Autoencoder (MAE) based on a Vision Transformer (ViT) architecture is utilized for pretraining. During this phase, a% of the image patches are masked, while the remaining (100 - a)% are processed by the encoder. The decoder block then reconstructs all patches, optimized via a customized loss function.*
+
+---
+
+## Getting Started
+
+### Prerequisites
+* Linux or macOS
+* Python 3.11+
+* NVIDIA GPU + CUDA toolkit (Recommended for training)
+
+### Environment Setup
+We recommend using `mamba` to manage dependencies. 
+
+> **Important Hardware Note:** > The `sdofmv2_environment.yml` file is configured for **CUDA 12.8** by default. If your hardware or drivers require a different CUDA version (e.g., CUDA 11.8), please open `sdofmv2_environment.yml` and modify the `pip` section at the bottom to match your system (e.g., change `cu128` to `cu118`) before running the setup commands.
+
+**Using Mamba:**
+```bash
+# Clone the repository
+git clone [https://github.com/Joaggi/sdofmv2.git](https://github.com/Joaggi/sdofmv2.git)
+cd sdofmv2
+
+# Create and activate the environment (This automatically installs PyTorch and the local package)
+mamba env create -f sdofmv2_environment.yml
+mamba activate sdofmv2
+```
+---
+
+## Repository Structure
+
+```text
+.
+├── configs/                # YAML configurations for experiments
+│   ├── downstream/         # Configs for downstream tasks (F10.7, solar wind)
+│   └── pretrain/           # Configs for MAE pretraining (AIA, HMI)
+├── notebooks/              # Jupyter notebooks for analysis and visualization
+│   ├── analysis/           # Attention maps, PCA, and masking analysis
+│   └── downstream_apps/    # How to use downstream scripts (Notebooks) for F10.7 and missing data applications
+├── scripts/                # Executable scripts for training and testing
+│   ├── pretrain.py         # Main pretraining script
+│   ├── finetuning_*.py     # Scripts for downstream finetuning
+│   └── test.py             # Script for evaluating checkpoints
+├── src/                    # Core source code package
+│   └── sdofmv2/
+│       ├── core/           # Base model architectures and modules
+│       ├── tasks/          # PyTorch Lightning modules (model & data module) for downstream tasks
+│       └── utils/          # Helper functions, physical constants and metrics
+├── pyproject.toml          # Project metadata and build dependencies
+└── sdofmv2_environment.yml # Mamba environment definition file
+```
+
+---
+
+## How to Use
+
+*(Note: It is recommended to run all scripts from the root directory of the repository so that file paths to `configs/` and `src/` resolve correctly.)*
+
+### 1. Data Preparation
+Before training or running inference, you need to prepare the dataset. 
+[Explain where to download the data, or provide a command if you have a script for it.]
+```bash
+python scripts/download_data_cache.py --target_dir ./assets/
+```
+
+### 2. Training the Model
+To train the model from scratch, execute the pretraining script and pass the relevant configuration file.
+```bash
+python scripts/pretrain.py --config-name pretrain_mae_AIA.yaml
+```
+
+### 3. Inference and Evaluation
+To evaluate a pre-trained checkpoint on the test set:
+```bash
+python scripts/test.py --config-name pretrain_mae_AIA.yaml
+```
+
+### 4. Downstream Finetuning
+To finetune the model on a specific downstream task (e.g., solar wind forecasting):
+```bash
+python scripts/finetuning_solarwind.py --config-name finetune_solarwind_config.yaml
+```
+---
+
+## Results & Visualizations
+[Include a brief summary of the model's performance. You can add a table of metrics or a sample plot showing predictions vs. ground truth.]
+
+![Sample Visualization](notebooks/analysis/SDOFMv2_AIA_results_exp.png) 
+*The first row displays the original ground-truth images. The second and third rows show the model's reconstructed images using masking ratios of 0% and 50%, respectively.*
+
+---
+
+## Citation
+If you find this repository or model useful in your academic research, please consider citing our work:
+
+```bibtex
+@misc{sdofmv2,
+  author = {Hong, Jinsu and Martin, Daniela and Gallego, Joseph},
+  title = {SDOFMv2: A Multi-Instrument Foundation Model for the Solar Dynamics Observatory with Transferable Downstream Applications},
+  year = {2026},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{[https://github.com/Joaggi/sdofmv2](https://github.com/Joaggi/sdofmv2)}},
+  note = {Jinsu Hong, Daniela Martin, and Joseph Gallego contributed equally to this work}
+}
+```
+
+## Contributing
+Contributions, bug reports, and feature requests are welcome! Please feel free to check the [issues page](https://github.com/Joaggi/sdofmv2/issues) or submit a pull request.

sdofmv2-0.1.0/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "sdofmv2"
+version = "0.1.0"
+description = "Solar phenomena prediction models"
+readme = "README.md"
+requires-python = ">=3.11"
+license = {text = "MIT"}
+authors = [
+    {name = "Joseph Gallego",     email = "joaggi@gmail.com"},
+    {name = "Daniela Martin",     email = "dmartinvega@gmail.com"},
+    {name = "Jinsu Hong",         email = "jinsuhong.knight@gmail.com"},
+]
+keywords = ["foundation model", "solar physics", "deep learning", "space weather"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Astronomy",
+]
+
+dependencies = [
+    "torch>=2.5.1",
+    "lightning>=2.6.0",
+    "numpy>=2.3.5",
+    "pandas>=2.3.3",
+    "transformers>=4.57.3",
+    "sunpy>=7.0.4",
+    "astropy>=6.0",
+    "timm>=1.0.22",
+    "einops>=0.8.1",
+    "hydra-core>=1.3.2",
+    "wandb>=0.23.1",
+]
+
+[project.urls]
+Repository = "https://github.com/Joaggi/sdofmv2"
+Issues = "https://github.com/Joaggi/sdofmv2/issues"
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

sdofmv2-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

sdofmv2-0.1.0/src/sdofmv2/core/__init__.py

@@ -0,0 +1,18 @@
+from .basemodule import BaseModule
+from .datamodule import (
+    SDOMLDataModule,
+    SDOMLDataset,
+    inverse_log_norm,
+    inverse_zscore_norm,
+)
+from .losses import (
+    mae_loss,
+    vector_aware_loss,
+    pixel_weight_loss,
+)
+from .mae3d import MaskedAutoencoderViT3D
+from .mae3d_old import MaskedAutoencoderViT3D_old
+from .mae_module import MAE
+from .mae_module_old import MAE_old
+from .pca_analysis import mapping_dense_to_rgb
+from .attention_map import plot_heads, patch_attn_layers, visualize_head

sdofmv2-0.1.0/src/sdofmv2/core/attention_map.py

@@ -0,0 +1,314 @@
+import os
+import types
+from typing import List, Optional
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import TwoSlopeNorm
+import sunpy.visualization.colormaps as sunpycm
+from sunpy.visualization.colormaps import color_tables
+
+import torch
+from loguru import logger as lgr_logger
+from omegaconf import OmegaConf
+from timm.layers import maybe_add_mask
+
+from sdofmv2.core import MAE, SDOMLDataModule
+
+
+# Attention patching function
+def patch_attn_layers(model: MAE) -> List[torch.Tensor]:
+    """
+    Monkey-patch the attention layers of a MAE model to store attention maps.
+
+    Returns:
+        attn_maps: List[Tensor] with shape [B, num_heads, N, N] per block
+    """
+    attn_maps: List[torch.Tensor] = []
+
+    def patched_forward(self, x, attn_mask=None):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        # Force unfused attention
+        if getattr(self, "fused_attn", False):
+            self.fused_attn = False
+
+        # Compute attention
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+        attn = maybe_add_mask(attn, attn_mask)
+        attn = attn.softmax(dim=-1)
+        attn_maps.append(attn.detach().cpu())
+        attn = self.attn_drop(attn)
+        x_out = attn @ v
+
+        x_out = x_out.transpose(1, 2).reshape(B, N, C)
+        x_out = self.norm(x_out)
+        x_out = self.proj(x_out)
+        x_out = self.proj_drop(x_out)
+        return x_out
+
+    # Patch all encoder blocks
+    for blk in model.autoencoder.blocks:
+        blk.attn.forward = types.MethodType(patched_forward, blk.attn)
+
+    return attn_maps
+
+
+def patch_id_to_xy(patch_id, patch_size=16, grid=32):
+    row = patch_id // grid
+    col = patch_id % grid
+    x = col * patch_size
+    y = row * patch_size
+    return x, y
+
+
+def attn_to_image(attn_vector, visible_patch_ids, img_size=512, patch_size=16):
+    heatmap = np.zeros((img_size, img_size), dtype=np.float32)
+
+    # Normalize for visualization
+    attn_norm = attn_vector / attn_vector.max()
+
+    for w, patch_id in zip(attn_norm, visible_patch_ids):
+        x, y = patch_id_to_xy(patch_id, patch_size)
+        heatmap[y : y + patch_size, x : x + patch_size] = w
+
+    return heatmap
+
+
+# Load MAE weights
+def load_mae_weights(ckpt_path: str, masking_ratio: float = 0.5) -> MAE:
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    hparams = ckpt["hyper_parameters"]
+
+    # if masking_ratio is not None:
+    #     hparams["masking_ratio"] = masking_ratio
+    #     print(f"Overriding masking_ratio to {masking_ratio}")
+
+    # Clean hyperparameters for model construction
+    for key in ["create_embedding_file", "lr", "num_classes"]:
+        if key in hparams:
+            hparams.pop(key)
+
+    if "wavelengths" in hparams:
+        hparams["chan_types"] = hparams.pop("wavelengths")
+
+    model = MAE(**hparams)
+    model.load_state_dict(ckpt["state_dict"], strict=False)
+    return model
+
+
+def visualize_head(attn_head, ids_keep, img_size=512, patch=16):
+    """
+    attn_head: [num_kept] attention values for ONLY the kept patches
+    ids_keep: [num_kept] indices of which patches were kept
+    """
+    grid_size = img_size // patch  # 32
+    heatmap = np.zeros((grid_size, grid_size))
+
+    # Normalize attention values
+    attn_norm = (attn_head - attn_head.min()) / (
+        attn_head.max() - attn_head.min() + 1e-8
+    )
+
+    # Place attention values at the correct patch positions
+    for score, patch_idx in zip(attn_norm, ids_keep):
+        row = patch_idx // grid_size
+        col = patch_idx % grid_size
+        heatmap[row, col] = score
+
+    # Resize to full image resolution
+    heatmap_full = np.repeat(np.repeat(heatmap, patch, axis=0), patch, axis=1)
+
+    return heatmap_full
+
+
+# def visualize_head(attn_head, ids_keep, img_size=512, patch=16):
+#     heatmap = np.zeros((img_size, img_size))
+
+#     # attn_head: [549] attention from one query token → each key token
+#     attn_norm = attn_head / (attn_head.max() + 1e-6)
+
+#     for score, patch_id in zip(attn_norm, ids_keep):
+#         y, x = divmod(patch_id.item(), 32)
+#         y *= patch
+#         x *= patch
+#         heatmap[y:y+patch, x:x+patch] = score
+
+#     return heatmap
+
+
+def plot_heads(
+    attn_maps, ids_restore, image, channels=["Bx", "By", "Bz"]
+):  # , patch_id=16
+    attn = attn_maps[0][0][:, 1:, 1:]  # [num_head, num_patch, num_patch]
+    num_heads = attn.shape[0]
+    full_order = torch.argsort(ids_restore[0])  # Invert argsort
+    num_kept = attn.shape[1]
+    ids_keep = full_order[:num_kept]
+
+    if len(channels) == 3:
+        cmap = [
+            sunpycm.cmlist.get("hmimag"),
+            sunpycm.cmlist.get("hmimag"),
+            sunpycm.cmlist.get("hmimag"),
+        ]
+        norm = TwoSlopeNorm(vmin=-4000, vcenter=0, vmax=4000)
+
+    elif len(channels) == 9:
+        cmap = [
+            sunpycm.cmlist.get("sdoaia131"),
+            sunpycm.cmlist.get("sdoaia1600"),
+            sunpycm.cmlist.get("sdoaia1700"),
+            sunpycm.cmlist.get("sdoaia171"),
+            sunpycm.cmlist.get("sdoaia193"),
+            sunpycm.cmlist.get("sdoaia211"),
+            sunpycm.cmlist.get("sdoaia304"),
+            sunpycm.cmlist.get("sdoaia335"),
+            sunpycm.cmlist.get("sdoaia94"),
+        ]
+        norm = None
+
+    else:
+        raise ValueError(f"Channel info is wrong")
+
+    attn_received = attn.mean(axis=1)
+    # attn_received = attn[:, 0, 1:]  # cls token
+    num_images = image.shape[0]
+    num_channels = image.shape[1]
+
+    fig, axs = plt.subplots(
+        num_images, num_heads + num_channels, figsize=(25, 4), squeeze=False
+    )
+
+    for i in range(image.shape[0]):
+        for i_ch, ch in enumerate(channels):
+            axs[i, i_ch].imshow(image[i, i_ch, :, :], cmap=cmap[i_ch], norm=norm)
+            axs[i, i_ch].set_title(f"{ch}")
+            axs[i, i_ch].axis("off")
+
+        for h in range(num_heads):
+            # head_attn = attn_shuffled[h, patch_id, :] # weights based on patches
+            head_attn = attn_received[h, :]
+            heatmap = visualize_head(head_attn, ids_keep, 512, 16)
+            axs[i, h + num_channels].imshow(heatmap, cmap="jet")
+            axs[i, h + num_channels].set_title(f"Head {h}")
+            axs[i, h + num_channels].axis("off")
+
+    plt.tight_layout()
+    # plt.savefig("attention_map_no_limb.png", dpi=200)
+    return fig, axs
+
+
+def plot_heads_no_limb(attn_maps, ids_restore, image, patch_id=16):
+    attn = attn_maps[0][0][:, 1:, 1:]  # [num_head, num_patch, num_patch]
+    num_heads = attn.shape[0]
+    ids_keep = ids_restore[0, : attn.shape[1]]
+    attn_shuffled = attn[:, ids_restore[0]][:, :, ids_restore[0]]
+    attn_received = attn_shuffled.mean(axis=1)
+    num_images = image.shape[0]
+    num_channels = image.shape[1]
+
+    fig, axs = plt.subplots(
+        num_images, num_heads + num_channels, figsize=(25, 4), squeeze=False
+    )
+
+    for i in range(image.shape[0]):
+        for i_ch, ch in enumerate(["Bx", "By", "Bz"]):
+            axs[i, i_ch].imshow(image[i, i_ch, :, :].to("cpu").numpy(), cmap="gray")
+            axs[i, i_ch].set_title(f"Ch: {ch}")
+            axs[i, i_ch].axis("off")
+
+        for h in range(num_heads):
+            # head_attn = attn_shuffled[h, patch_id, :] # weights based on patches
+            head_attn = attn_received[h, :]  # weights based on mean of patches
+            heatmap = visualize_head(head_attn, ids_keep, 512, 16)
+            axs[i, h + num_channels].imshow(heatmap, cmap="jet")
+            axs[i, h + num_channels].set_title(f"Head {h}")
+            axs[i, h + num_channels].axis("off")
+
+    plt.tight_layout()
+    # plt.savefig("attention_map_no_limb.png", dpi=200)
+    return fig
+
+
+# Main execution
+if __name__ == "__main__":
+    cfg = OmegaConf.load(
+        "/home/jhong36/Project/2025-HL-Solar-Wind/solar_phenomena_prediction/configs/pretrain_mae.yaml"
+    )
+
+    # Setup dataset
+    data_module = SDOMLDataModule(
+        hmi_path=(
+            os.path.join(
+                cfg.data.sdoml.base_directory, cfg.data.sdoml.sub_directory.hmi
+            )
+            if cfg.data.sdoml.sub_directory.hmi
+            else None
+        ),
+        aia_path=(
+            os.path.join(
+                cfg.data.sdoml.base_directory, cfg.data.sdoml.sub_directory.aia
+            )
+            if cfg.data.sdoml.sub_directory.aia
+            else None
+        ),
+        eve_path=None,
+        components=cfg.data.sdoml.components,
+        wavelengths=cfg.data.sdoml.wavelengths,
+        ions=cfg.data.sdoml.ions,
+        frequency=cfg.data.sdoml.frequency,
+        batch_size=cfg.model.opt.batch_size,
+        num_workers=cfg.data.num_workers,
+        val_months=cfg.data.month_splits.val,
+        test_months=cfg.data.month_splits.test,
+        holdout_months=cfg.data.month_splits.holdout,
+        cache_dir=os.path.join(
+            cfg.data.sdoml.save_directory, cfg.data.sdoml.sub_directory.cache
+        ),
+        min_date=cfg.data.min_date,
+        max_date=cfg.data.max_date,
+        num_frames=cfg.model.mae.num_frames,
+        drop_frame_dim=cfg.data.drop_frame_dim,
+        apply_mask=cfg.data.sdoml.apply_mask,
+        precision=cfg.experiment.precision,
+        normalization=cfg.data.sdoml.normalization,
+    )
+    data_module.setup()
+
+    # Load model
+    base_path = "../../../../assets/check_point/backbone/"
+    model_hmi = load_mae_weights(
+        os.path.join(base_path, "id_xn2c11go_mae_epoch=25-val_loss=0.00.ckpt"),
+        # masking_ratio=0  # full image by default
+    )
+    # model_hmi.autoencoder.ids_limb_mask = None
+    model_hmi.eval()
+
+    # Patch attention layers
+    attn_maps = patch_attn_layers(model_hmi)
+
+    # Forward pass
+    id_input = 0
+    x = data_module.test_ds[id_input][0].unsqueeze(0)
+    lgr_logger.info(f"Input shape: {x.shape}")
+
+    with torch.no_grad():
+        latent, mask, ids_restore = model_hmi.autoencoder.forward_encoder(
+            x, mask_ratio=0
+        )
+
+    # Example: check first attention map
+    if len(attn_maps) > 0:
+        lgr_logger.info(f"First attention map shape: {attn_maps[0].shape}")
+
+    fig = plot_heads(attn_maps, ids_restore, x[:, 0:1, 0, :, :])
+    # fig = plot_heads_no_limb(attn_maps, ids_restore, x[:, 0:1, 0, :, :])