🔮 CHGNet Fine-tuning Results

# ============================================================
# Fine-tuning CHGNet on a Subset of MP20 Dataset with Hyperparameter Sweep
# ============================================================

import os
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend for thread safety
import matplotlib.pyplot as plt
from tqdm import tqdm
from datasets import Dataset
from pymatgen.core import Structure
from pymatgen.ext.matproj import MPRester
from chgnet.data.dataset import StructureData, get_train_val_test_loader
from chgnet.model import CHGNet
from chgnet.trainer import Trainer
from sklearn.metrics import r2_score
import torch
import wandb
from pathlib import Path

# ============================================================
# Configuration
# ============================================================

MP_API_KEY = os.getenv("MP_API_KEY")  # Set as environment variable
DATASET_PATH = "mp_20.json"  # Path to your MP20 dataset
N_TEST_SAMPLES = 1000
OUTPUT_DIR = Path("sweep_results")
OUTPUT_DIR.mkdir(exist_ok=True)

# Sweep configuration
SWEEP_CONFIG = {
    "method": "grid",
    "metric": {"name": "val_mae", "goal": "minimize"},
    "parameters": {
        "n_train_samples": {"values": [100, 250, 500, 1000, 1500, 2000]},
        "batch_size": {"values": [8, 16]},
        "learning_rate": {"values": [1e-3, 5e-3, 1e-2]},
        "epochs": {"values": [5, 10]},
    }
}

# ============================================================
# Utility Functions
# ============================================================

def get_structure_from_mp_id(mp_id: str) -> Structure | None:
    """Fetches structure from Materials Project using MPRester."""
    try:
        with MPRester(api_key=MP_API_KEY) as mpr:
            return mpr.get_structure_by_material_id(mp_id)
    except Exception as e:
        print(f"❌ Could not fetch structure for {mp_id}: {e}")
        return None

def energy_per_atom_fn(model):
    """Returns a closure to compute predicted energy per atom."""
    def _energy_per_atom(structure):
        try:
            return model.predict_structure(structure)["e"]
        except Exception as e:
            return np.nan
    return _energy_per_atom

# ... [Full script continues - see GitHub repository]

# ============================================================
# Main Execution
# ============================================================

if __name__ == "__main__":
    wandb.login()
    
    # Prepare data once
    dataset = Dataset.from_json(DATASET_PATH)
    base_model = CHGNet.load()
    
    MAX_SAMPLES = 2000
    PREPARED_STRUCTURES, PREPARED_ENERGIES, PREPARED_FORCES = prepare_max_training_data(
        dataset, MAX_SAMPLES, base_model
    )
    
    EVAL_STRUCTURES, EVAL_GROUND_TRUTH = prepare_eval_data(dataset, N_TEST_SAMPLES)
    
    # Create and run sweep
    sweep_id = wandb.sweep(SWEEP_CONFIG, project="chgnet-finetuning-sweep")
    wandb.agent(sweep_id, train_with_config, count=None)
    
    print("🎉 Sweep completed!")