load_data.py

from astropy.io import fits
import numpy as np
import os
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from astropy.convolution import Box2DKernel
from scipy.signal import convolve2d
import hashlib
from torch.utils.data import DataLoader, SubsetRandomSampler

class StarPSFDatasetFromFITS(Dataset):
    def __init__(self, fits_path, iccd_lst, device):
        self.fits_path = fits_path
        self.device = device
        self.iccd_lst = iccd_lst
        self.data = {}
        self.indices = []

        self._load_data()
    def _load_data(self):
        with fits.open(self.fits_path) as hdul:
            for iccd in self.iccd_lst:
                star_data = hdul[f"STARS_{iccd:02d}"].data
                psf_data = hdul[f"PSFS_{iccd:02d}"].data
                
                meta = hdul[f"METADATA_{iccd:02d}"].data
                obj_ids = meta['OBJ_ID']
                self.data[iccd] = {
                    'stars': star_data,
                    'psfs': psf_data,
                    'meta': meta,
                    'obj_ids': obj_ids
                }
                for i in range(len(obj_ids)):
                    self.indices.append((iccd, i))

    def __len__(self):
        return len(self.indices)

    def max_normalize(self, array):
        minv, maxv = np.min(array), np.max(array)
        if maxv > minv:
            return ((array - minv) / (maxv - minv)).astype(np.float32)
        return array.astype(np.float32)
    
    def sum_normalize(self, array):
        return (array/np.sum(array)).astype(np.float32)

    def __getitem__(self, idx):
        iccd, i = self.indices[idx]
        star = self.sum_normalize(self.data[iccd]['stars'][i])
        psf = convolve2d(self.data[iccd]['psfs'][i], Box2DKernel(width=2), mode='same')
        meta = self.data[iccd]['meta'][i]

        return {
            'lr_star': transforms.ToTensor()(star).to(self.device),
            'hr_psf': transforms.ToTensor()(psf).to(self.device),
            'obj_id': int(meta['OBJ_ID']),
            'iccd': int(meta['ICCD']),
            'x_image': float(meta['X_IMAGE']),
            'y_image': float(meta['Y_IMAGE']),
        }

class DataSplitter:
    def __init__(self, dataset, save_dir=None, seed=42):
        self.dataset = dataset
        self.save_dir = save_dir
        if save_dir is not None:
            os.makedirs(save_dir, exist_ok=True)
        self.seed = seed 
        self.data_hash = self._generate_data_hash()
        
        self.train_loader = None
        self.val_loader = None
        self.test_loader = None
        
    def _generate_data_hash(self):
        hash_obj = hashlib.md5()
        for iccd in self.dataset.iccd_lst:
            for obj_id in self.dataset.data[iccd]['obj_ids']:
                hash_obj.update(str(obj_id).encode())
        return hash_obj.hexdigest()

    def _get_collate_fn(self):
        return lambda batch: {
            'lr_star': torch.stack([x['lr_star'] for x in batch]),
            'hr_psf': torch.stack([x['hr_psf'] for x in batch]),
            'psfex_psf': torch.stack([x['psfex_psf'] for x in batch]) if 'psfex_psf' in batch[0] else None,
            'iccd': [x['iccd'] for x in batch],
            'obj_id': [x['obj_id'] for x in batch],
            'x_image': [x['x_image'] for x in batch],
            'y_image': [x['y_image'] for x in batch],
        }

    def create_train_val_loaders(self, batch_size=16, val_ratio=0.2, max_batches=None):
        collate_fn = self._get_collate_fn()
        dataset_size = len(self.dataset)
        indices = list(range(dataset_size))
        val_split = int(np.floor(val_ratio * dataset_size))
        np.random.seed(self.seed)
        np.random.shuffle(indices)    
        train_indices, val_indices = indices[val_split:], indices[:val_split] 
        if max_batches is not None:
            max_train_samples = max_batches * batch_size
            train_indices = train_indices[:max_train_samples]        
            
        self.train_loader = DataLoader(
            self.dataset, batch_size=batch_size,
            sampler=SubsetRandomSampler(train_indices),
            collate_fn=collate_fn)
    
        self.val_loader = DataLoader(
            self.dataset, batch_size=batch_size,
            sampler=SubsetRandomSampler(val_indices),
            collate_fn=collate_fn)
        
        self._save_train_val_stats()
        
        return self.train_loader, self.val_loader

    def create_test_loader(self, batch_size=16):
        collate_fn = self._get_collate_fn()
        
        self.test_loader = DataLoader(
            self.dataset, batch_size=batch_size,
            shuffle=False,  
            collate_fn=collate_fn)
        
        self._save_test_stats()
        return self.test_loader

    def _save_train_val_stats(self):
        if self.save_dir is None:
            return
            
        train_counts = self._count_samples_by_iccd(self.train_loader)
        val_counts = self._count_samples_by_iccd(self.val_loader)
        
        file_path = os.path.join(self.save_dir, "train_val_stats.txt")
        with open(file_path, 'w') as f:
            f.write("#iccd train_count val_count\n")
            for iccd in sorted(train_counts.keys()):
                train_num = train_counts.get(iccd, 0)
                val_num = val_counts.get(iccd, 0)
                f.write(f"{iccd} {train_num} {val_num}\n")
        
        print(f"Train/validation statistics saved to {file_path}")

    def _save_test_stats(self):
        if self.save_dir is None:
            return
            
        test_counts = self._count_samples_by_iccd(self.test_loader)
        file_path = os.path.join(self.save_dir, "test_stats.txt")
        with open(file_path, 'w') as f:
            f.write("#iccd test_count\n")
            for iccd in sorted(test_counts.keys()):
                test_num = test_counts.get(iccd, 0)
                f.write(f"{iccd} {test_num}\n")
        

    def _count_samples_by_iccd(self, dataloader):
        """统计数据加载器中每个ICCD的样本数量"""
        if dataloader is None:
            return {}
            
        counts = {}
        for batch in dataloader:
            for iccd in batch['iccd']:
                counts[iccd] = counts.get(iccd, 0) + 1
        return counts

class Finetune_StarPSFDataset(Dataset):
    def __init__(self, fits_path, psfex_psf_path, iccd_lst, device):
        self.fits_path = fits_path
        self.psfex_psf_path = psfex_psf_path
        self.device = device
        self.iccd_lst = iccd_lst
        self.data = {}
        self.indices = []
        self.psfex_index = {}  

        self._build_psfex_index()  
        self._load_data()          

    def _build_psfex_index(self):
        for iccd in self.iccd_lst:
            psfex_file = os.path.join(self.psfex_psf_path, f"psf_ccd{iccd}.fits")
            mapping = {}
            if os.path.exists(psfex_file):
                with fits.open(psfex_file) as hdul:
                    meta_table = hdul['METADATA'].data
                    obj_ids = np.asarray(meta_table['OBJ_ID']).astype(np.int64)
                    mapping = {int(oid): int(i) for i, oid in enumerate(obj_ids)}
            else:
                print(f"[WARN] PSFEx file not found: {psfex_file}. All samples on CCD {iccd} will be skipped.")
            self.psfex_index[iccd] = mapping

    def _load_data(self):
        skipped = 0
        with fits.open(self.fits_path) as hdul:
            for iccd in self.iccd_lst:
                star_data = hdul[f"STARS_{iccd:02d}"].data
                psf_data = hdul[f"PSFS_{iccd:02d}"].data
                meta = hdul[f"METADATA_{iccd:02d}"].data
                obj_ids = np.asarray(meta['OBJ_ID']).astype(np.int64)

                self.data[iccd] = {
                    'stars': star_data,
                    'psfs': psf_data,
                    'meta': meta,
                    'obj_ids': obj_ids
                }

                available = self.psfex_index.get(iccd, {})
                for i, oid in enumerate(obj_ids):
                    if int(oid) in available:
                        self.indices.append((iccd, i))
                    else:
                        skipped += 1

        if skipped > 0:
            print(f"[INFO] Skipped {skipped} samples without PSFEx match.")

    def _load_psfex_psf(self, iccd, obj_id):
        psfex_file = os.path.join(self.psfex_psf_path, f"psf_ccd{iccd}.fits")
        idx = self.psfex_index.get(iccd, {}).get(int(obj_id), None)
        if idx is None or not os.path.exists(psfex_file):
            return None, None

        with fits.open(psfex_file) as hdul:
            psf_data = hdul["RECON_PSF"].data
            meta_table = hdul['METADATA'].data
            return psf_data[idx], meta_table[idx]

    def __len__(self):
        return len(self.indices)

    def max_normalize(self, array):
        minv, maxv = np.min(array), np.max(array)
        if maxv > minv:
            return ((array - minv) / (maxv - minv)).astype(np.float32)
        return array.astype(np.float32)
    
    def sum_normalize(self, array):
        return (array/np.sum(array)).astype(np.float32)

    @staticmethod
    def _to_native_float32(arr: np.ndarray) -> np.ndarray:
        arr = np.asarray(arr, dtype=np.float32)
        if arr.dtype.byteorder not in ('=', '|'):
            arr = arr.byteswap().newbyteorder()
        return np.ascontiguousarray(arr)
    
    def __getitem__(self, idx):
        iccd, i = self.indices[idx]
        star = self.sum_normalize(self.data[iccd]['stars'][i])
        psf = convolve2d(self.data[iccd]['psfs'][i], Box2DKernel(width=2), mode='same')
        meta = self.data[iccd]['meta'][i]

        psfex_psf, psfex_meta = self._load_psfex_psf(int(meta['ICCD']), int(meta['OBJ_ID']))
        if psfex_psf is None:
            raise KeyError(f"PSFEx match not found for ICCD={meta['ICCD']} OBJ_ID={meta['OBJ_ID']}")

        psfex_psf = self._to_native_float32(psfex_psf)

        return {
            'lr_star': transforms.ToTensor()(star).to(self.device),
            'hr_psf': transforms.ToTensor()(psf).to(self.device),
            'psfex_psf': transforms.ToTensor()(psfex_psf).to(self.device),
            'obj_id': int(meta['OBJ_ID']),
            'iccd': int(meta['ICCD']),
            'x_image': float(meta['X_IMAGE']),
            'y_image': float(meta['Y_IMAGE']),
            'x_focal': float(meta['XFocal']),
            'y_focal': float(meta['YFocal'])
        }