data_module.py

"""
Data Module for Li-Si-P-S Amorphous Structure Generation

This module provides:
- AmorphousStructureDataset: Dataset for loading atomic structures
- StructureDataModule: Lightning DataModule wrapper

Designed for unconditional generation (no target properties like XANES).
"""

from typing import Optional
import torch
import numpy as np
import lightning.pytorch as pl

import ase.io
from pathlib import Path
from torch_geometric.data import Data, Dataset, DataLoader
from sklearn.preprocessing import OneHotEncoder
from scipy.spatial.transform import Rotation

from graphite import periodic_radius_graph, VarianceExplodingDiffuser


# Li-Si-P-S atomic numbers for reference
ATOMIC_NUMBERS = {
    'Li': 3,
    'Si': 14,
    'P': 15,
    'S': 16,
}


class AmorphousStructureDataset(Dataset):
    """
    PyTorch Geometric Dataset for amorphous Li-Si-P-S structures.
    
    This dataset:
    1. Loads structures from VASP/CIF/XYZ files
    2. Creates atomic graphs using periodic radius graph
    3. Applies forward diffusion noise for score matching
    4. Applies random rotation for data augmentation
    
    Args:
        data_dir: Path to directory containing structure files
        file_pattern: Glob pattern for structure files (default: '*.vasp')
        cutoff: Cutoff radius for graph construction (Angstroms)
        k: Maximum noise level for variance exploding diffusion
        dup: Number of times to duplicate data for longer epochs
    """
    def __init__(self, data_dir: str, file_pattern: str = '*.vasp', 
                 cutoff: float = 3.0, k: float = 0.8, dup: int = 1):
        super().__init__()
        self.data_dir = data_dir
        self.cutoff = cutoff
        self.diffuser = VarianceExplodingDiffuser(k=k)
        self.dup = dup

        # Find all structure files
        self.structure_files = sorted(Path(data_dir).glob(file_pattern))
        
        if len(self.structure_files) == 0:
            raise FileNotFoundError(f"No files matching '{file_pattern}' found in {data_dir}")
        
        print(f"Found {len(self.structure_files)} structure files")

        # Load all structures
        atoms_list = []
        for f in self.structure_files:
            try:
                atoms = ase.io.read(f)
                atoms.wrap()
                atoms_list.append(atoms)
            except Exception as e:
                print(f"Warning: Could not read {f}: {e}")
        
        if len(atoms_list) == 0:
            raise ValueError("No valid structures loaded!")
        
        print(f"Loaded {len(atoms_list)} structures")

        # Create one-hot encoder for atomic species
        # Collect all unique atomic numbers across all structures
        all_numbers = np.concatenate([atoms.numbers for atoms in atoms_list])
        unique_numbers = np.unique(all_numbers)
        
        self.atom_encoder = OneHotEncoder(sparse_output=False)
        self.atom_encoder.fit(unique_numbers.reshape(-1, 1))
        
        self.num_species = len(unique_numbers)
        print(f"Number of species: {self.num_species}")
        print(f"Species (atomic numbers): {unique_numbers}")

        # Convert to PyG Data objects
        self.dataset = []
        for atoms in atoms_list:
            z = self.atom_encoder.transform(atoms.numbers.reshape(-1, 1))
            data = Data(
                z    = torch.tensor(z, dtype=torch.float),
                pos  = torch.tensor(atoms.positions, dtype=torch.float),
                cell = np.array(atoms.cell),
            )
            
            # Duplicate for longer epochs
            for _ in range(dup):
                self.dataset.append(data.clone())
        
        print(f"Dataset size (with duplication): {len(self.dataset)}")
    
    def len(self):
        return len(self.dataset)

    def get(self, idx):
        """Get a sample with diffusion noise and graph construction."""
        data = self.dataset[idx].clone()
        data = self._diffuse_pos(data)
        data = self._build_graph(data)
        data = self._random_rotate(data)
        return data
        
    def _diffuse_pos(self, data):
        """Apply forward diffusion noise to positions."""
        t = torch.rand(1).clip(self.diffuser.t_min, self.diffuser.t_max)
        data.t = t.expand(data.pos.size(0), 1)
        data.pos, data.eps_r = self.diffuser.forward_noise(data.pos, data.t)
        data.sigma_r = self.diffuser.sigma(data.t)
        return data
    
    def _build_graph(self, data):
        """Construct periodic radius graph."""
        cell = torch.tensor(data.cell, dtype=torch.float)
        data.edge_index, edge_vec = periodic_radius_graph(data.pos, self.cutoff, cell=cell)
        data.edge_len = edge_vec.norm(dim=-1, keepdim=True)
        data.edge_attr = torch.hstack([edge_vec, data.edge_len])
        return data

    def _random_rotate(self, data):
        """Apply random 3D rotation for augmentation."""
        R = torch.tensor(Rotation.random().as_matrix(), dtype=torch.float)
        data.pos = data.pos @ R
        data.edge_attr[:, :3] = data.edge_attr[:, :3] @ R
        data.eps_r = data.eps_r @ R
        return data


class StructureDataModule(pl.LightningDataModule):
    """
    PyTorch Lightning DataModule for structure data.
    
    Args:
        data_dir: Path to directory containing structure files
        file_pattern: Glob pattern for files (default: '*.vasp')
        cutoff: Cutoff radius for graph construction
        k: Maximum noise level for diffusion
        dup: Data duplication factor
        batch_size: Graphs per batch
        num_workers: DataLoader workers
    """
    def __init__(self, data_dir: str, file_pattern: str = '*.vasp',
                 cutoff: float = 3.0, k: float = 0.8, dup: int = 1,
                 batch_size: int = 8, num_workers: int = 4):
        super().__init__()
        self.save_hyperparameters()

        self.data_dir = data_dir
        self.file_pattern = file_pattern
        self.cutoff = cutoff
        self.k = k
        self.dup = dup
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        """Create dataset."""
        self.train_set = AmorphousStructureDataset(
            data_dir=self.data_dir,
            file_pattern=self.file_pattern,
            cutoff=self.cutoff,
            k=self.k,
            dup=self.dup,
        )

    def train_dataloader(self):
        return DataLoader(
            self.train_set, 
            shuffle=True,  
            batch_size=self.batch_size, 
            num_workers=self.num_workers, 
            pin_memory=True
        )
    
    def teardown(self, stage:Optional[str] = None):
        """Clean up after training."""
        pass
    
    @property
    def num_species(self):
        """Get number of atomic species (available after setup)."""
        if hasattr(self, 'train_set'):
            return self.train_set.num_species
        return None
    
    @property
    def diffuser(self):
        """Get diffuser instance (available after setup)."""
        if hasattr(self, 'train_set'):
            return self.train_set.diffuser
        return None