psf_shape_meas_plot.py

from astropy.io import fits
import matplotlib.pyplot as plt
from glob import glob
import os
import numpy as np
import pandas as pd
import seaborn as sns
import galsim
from tqdm import tqdm
from astropy.table import Table, vstack
from scipy.interpolate import griddata
from astropy.visualization import simple_norm

iccd_positions = {
        6: (0, "y1"), 7: (1, "i1"), 8: (2, "g1"), 9: (3, "r1"),
        11: (5, "z1"), 12: (6, "NUV1"), 13: (7, "NUV2"), 14: (8, "u1"),
        15: (9, "y2"), 16: (10, "y3"), 17: (11, "u2"), 18: (12, "NUV3"),
        19: (13, "NUV4"), 20: (14, "z2"), 22: (16, "r2"),
        23: (17, "g2"), 24: (18, "i2"), 25: (19, "y4")
    }
filter_CenWave = {
    "NUV": 2888, "u": 3592, "g": 4820, "r": 6181, "i": 7640, "z": 8955, "y": 9584,
    "GU": 3310, "GV": 5248, "GI": 7770
}
def get_subplot_pos(iccd):
    """为每个CCD分配固定的子图位置"""
    pos_map = {
        6:1, 7:2, 8:3, 9:4,
        11:6, 12:7, 13:8, 14:9, 15:10,
        16:11, 17:12, 18:13, 19:14, 20:15, 
        22:17, 23:18, 24:19, 25:20
    }
    return pos_map.get(iccd, 1)

def plot_metrics_from_txt(txt_file_path, save_dir=None):
    df = pd.read_csv(txt_file_path, delim_whitespace=True)
    ccd_groups = {}
    for iccd, group in df.groupby('CCD_ID'):
        ccd_groups[iccd] = {
            'obj_id': group['OBJ_ID'].values,
            'x_pos': group['X_IMAGE'].values,
            'y_pos': group['Y_IMAGE'].values,
            'all_psnr': group['PSNR(dB)'].values,
            'all_ssim': group['SSIM'].values,
            'all_rmse': group['RMSE'].values
        }
    plot_histograms(ccd_groups, save_dir)
    plot_spatial_distributions(ccd_groups, save_dir)
    
def plot_histograms(ccd_groups, save_dir=None):
    """绘制PSNR/SSIM/RMSE的直方图（4行5列布局）"""
    # 创建三个大图（PSNR/SSIM/RMSE）
    fig_psnr, axs_psnr = plt.subplots(4, 5, figsize=(24, 18))
    fig_ssim, axs_ssim = plt.subplots(4, 5, figsize=(24, 18))
    fig_rmse, axs_rmse = plt.subplots(4, 5, figsize=(24, 18))
    
    for iccd, data in ccd_groups.items():
        pos = get_subplot_pos(iccd)
        row = (pos - 1) // 5
        col = (pos - 1) % 5
        # 计算统计量
        psnr_median, psnr_std = np.median(data['all_psnr']), np.std(data['all_psnr'])
        ssim_median, ssim_std = np.median(data['all_ssim']), np.std(data['all_ssim'])
        rmse_median, rmse_std = np.median(data['all_rmse']), np.std(data['all_rmse'])
        
        # 绘制PSNR直方图
        sns.histplot(data['all_psnr'], ax=axs_psnr[row, col], kde=True, color='skyblue', bins=20)
        axs_psnr[row, col].axvline(psnr_median, color='r', linestyle='--', linewidth=1.5)
        axs_psnr[row, col].set_title(f'CCD {iccd}\nMedian: {psnr_median:.2f} ± {psnr_std:.2f}')
        axs_psnr[row, col].tick_params(direction='in')
        
        # 绘制SSIM直方图
        sns.histplot(data['all_ssim'], ax=axs_ssim[row, col], kde=True, color='lightgreen', bins=20)
        axs_ssim[row, col].axvline(ssim_median, color='r', linestyle='--', linewidth=1.5)
        axs_ssim[row, col].set_title(f'CCD {iccd}\nMedian: {ssim_median:.6f} ± {ssim_std:.6f}')
        axs_ssim[row, col].tick_params(direction='in')
        # 绘制RMSE直方图
        sns.histplot(data['all_rmse'], ax=axs_rmse[row, col], kde=True, color='salmon', bins=20)
        axs_rmse[row, col].axvline(rmse_median, color='r', linestyle='--', linewidth=1.5)
        axs_rmse[row, col].set_title(f'CCD {iccd}\nMedian: {rmse_median:.6f} ± {rmse_std:.6f}')
        axs_rmse[row, col].tick_params(direction='in')
            
    # 调整布局并保存
    for fig, ax, name in zip([fig_psnr, fig_ssim, fig_rmse], [axs_psnr, axs_ssim, axs_rmse], ['psnr', 'ssim', 'rmse']):
        fig.suptitle(f'{name.upper()} Histograms', y=1.0, fontsize=18)
        fig.tight_layout()
        fig.delaxes(ax[0, 4]) 
        fig.delaxes(ax[3, 0]) 
        fig.savefig(os.path.join(save_dir, f'{name}_histograms_4x5.png'), dpi=300, bbox_inches='tight')
        plt.close(fig)

def plot_spatial_distributions(ccd_groups, save_dir=None):
    fig_psnr, axs_psnr = plt.subplots(4, 5, figsize=(24, 18))
    fig_ssim, axs_ssim = plt.subplots(4, 5, figsize=(24, 18))
    fig_rmse, axs_rmse = plt.subplots(4, 5, figsize=(24, 18))
    
    for iccd, data in ccd_groups.items():
        pos = get_subplot_pos(iccd)
        row = (pos - 1) // 5
        col = (pos - 1) % 5
        
        # PSNR空间分布
        sc1 = axs_psnr[row, col].scatter(data['x_pos'], data['y_pos'], c=data['all_psnr'], cmap='viridis', s=10)
        axs_psnr[row, col].set_title(f'CCD {iccd}')
        axs_psnr[row, col].set_xticks([])
        axs_psnr[row, col].set_yticks([])
        axs_psnr[row, col].invert_yaxis() 
        fig_psnr.colorbar(sc1, ax=axs_psnr[row, col])
        
        
        # SSIM空间分布
        sc2 = axs_ssim[row, col].scatter(data['x_pos'], data['y_pos'], c=data['all_ssim'], cmap='plasma', s=10)
        axs_ssim[row, col].set_title(f'CCD {iccd}')
        axs_ssim[row, col].axis('off')
        axs_ssim[row, col].set_xticks([])
        axs_ssim[row, col].set_yticks([])
        axs_ssim[row, col].invert_yaxis() 
        fig_ssim.colorbar(sc2, ax=axs_ssim[row, col])
        
        # RMSE空间分布
        sc3 = axs_rmse[row, col].scatter(data['x_pos'], data['y_pos'], c=data['all_rmse'], cmap='inferno', s=10)
        axs_rmse[row, col].set_title(f'CCD {iccd}')
        axs_rmse[row, col].set_xticks([])
        axs_rmse[row, col].set_yticks([])
        axs_rmse[row, col].invert_yaxis() 
        fig_rmse.colorbar(sc3, ax=axs_rmse[row, col])
    
    # 调整布局并保存
    for fig, ax, name in zip([fig_psnr, fig_ssim, fig_rmse], [axs_psnr, axs_ssim, axs_rmse], ['psnr', 'ssim', 'rmse']):
        fig.suptitle(f'{name.upper()} Spatial Distribution', y=1.0, fontsize=18)
        fig.tight_layout()
        fig.delaxes(ax[0, 4]) 
        fig.delaxes(ax[3, 0]) 
        fig.savefig(os.path.join(save_dir, f'{name}_spatial_4x5.png'), dpi=300, bbox_inches='tight')
        plt.close(fig)

def plot_compare_psnr_ssim_rmse(psfex_txt, model_txt, save_dir, mode = ' ', finetune_txt=''):
    def get_ccd_groups(df):
        ccd_groups = {}
        for iccd, group in df.groupby('CCD_ID'):
            ccd_groups[iccd] = {
                'obj_id': group['OBJ_ID'].values,
                'x_pos': group['X_IMAGE'].values,
                'y_pos': group['Y_IMAGE'].values,
                'psnr': group['PSNR(dB)'].values,
                'ssim': group['SSIM'].values,
                'rmse': group['RMSE'].values
            }
        return ccd_groups

    df1 = pd.read_csv(psfex_txt, delim_whitespace=True)
    df2 = pd.read_csv(model_txt, delim_whitespace=True)
    ccd_ids1 = df1['CCD_ID'].unique()
    ccd_ids2 = df2['CCD_ID'].unique()
    if not np.array_equal(ccd_ids1, ccd_ids2):
        raise ValueError("两个文件的CCD_ID不相同")

    ccd_groups1 = get_ccd_groups(df1)
    ccd_groups2 = get_ccd_groups(df2)

    bins = 30  

    fig_psnr, axs_psnr = plt.subplots(4, 5, figsize=(24, 15))
    fig_ssim, axs_ssim = plt.subplots(4, 5, figsize=(24, 15))
    fig_rmse, axs_rmse = plt.subplots(4, 5, figsize=(24, 15))
    
    if mode == 'finetune':
        df3 = pd.read_csv(finetune_txt, delim_whitespace=True)
        ccd_groups3 = get_ccd_groups(df3)
        
        
    for iccd in ccd_ids1:
        pos = get_subplot_pos(iccd)
        row = (pos - 1) // 5
        col = (pos - 1) % 5
        data1 = ccd_groups1[iccd]
        data2 = ccd_groups2[iccd]
        
        
        if mode == 'finetune':
            data3 = ccd_groups3[iccd]
            psnr_min = min(np.min(data1['psnr']), np.min(data2['psnr']), np.min(data3['psnr']))
            psnr_max = max(np.max(data1['psnr']), np.max(data2['psnr']), np.max(data3['psnr']))
            ssim_min = min(np.min(data1['ssim']), np.min(data2['ssim']), np.min(data3['ssim']))
            ssim_max = max(np.max(data1['ssim']), np.max(data2['ssim']), np.max(data3['ssim']))
            rmse_min = min(np.min(data1['rmse']), np.min(data2['rmse']), np.min(data3['rmse']))
            rmse_max = max(np.max(data1['rmse']), np.max(data2['rmse']), np.max(data3['rmse']))            
        else:
            psnr_min = min(np.min(data1['psnr']), np.min(data2['psnr']))
            psnr_max = max(np.max(data1['psnr']), np.max(data2['psnr']))
            ssim_min = min(np.min(data1['ssim']), np.min(data2['ssim']))
            ssim_max = max(np.max(data1['ssim']), np.max(data2['ssim']))
            rmse_min = min(np.min(data1['rmse']), np.min(data2['rmse']))
            rmse_max = max(np.max(data1['rmse']), np.max(data2['rmse']))        
            
            
        psnr_range = (psnr_min - 0.5, psnr_max + 0.5)  # 边界稍微留白
        sns.histplot(data1['psnr'], ax=axs_psnr[row, col], element='step', fill=True, alpha=0.8,
                     color='#4E79A7', bins=bins, binrange=psnr_range, label='PSFEx', linewidth=0)
        sns.histplot(data2['psnr'], ax=axs_psnr[row, col], element='step', fill=True, alpha=0.8,
                color='#79706E', bins=bins, binrange=psnr_range, label='Model', linewidth=0)
                    
        if mode == 'finetune':
            sns.histplot(data3['psnr'], ax=axs_psnr[row, col], element='step', fill=True, alpha=0.8,
                    color="#40747E", bins=bins, binrange=psnr_range, label='finetune', linewidth=0)
           
        axs_psnr[row, col].set_title(f'CCD {iccd}', fontsize=16)
        axs_psnr[row, col].set_xlim(psnr_range)
        axs_psnr[row, col].legend(loc='upper left', fontsize=16)
        axs_psnr[row, col].tick_params(direction='in')
        
        # SSIM
        ssim_range = (ssim_min - 0.001, ssim_max + 0.001)
        sns.histplot(data1['ssim'], ax=axs_ssim[row, col], color='#8CD17D', bins=bins, binrange=ssim_range,
                     element='step', fill=True, alpha=0.7, linewidth=0.2, label='PSFEx')
        sns.histplot(data2['ssim'], ax=axs_ssim[row, col], color='#D7B5A6', bins=bins, binrange=ssim_range,
                    element='step', fill=True, alpha=0.7, linewidth=0.2, label='Model')      
           
        if mode == 'finetune':
            sns.histplot(data3['ssim'], ax=axs_ssim[row, col], color="#AC7366", bins=bins, binrange=ssim_range,
                        element='step', fill=True, alpha=0.7, linewidth=0.2, label='finetune')

                    
        axs_ssim[row, col].set_title(f'CCD {iccd}', fontsize=16)
        axs_ssim[row, col].set_xlim(ssim_range)
        axs_ssim[row, col].legend(loc='upper left', fontsize=16)
        axs_ssim[row, col].tick_params(direction='in')
        
        # RMSE
        rmse_range = (rmse_min - 0.001, rmse_max + 0.001)
        sns.histplot(data1['rmse'], ax=axs_rmse[row, col], color='#edf8b1', bins=bins, binrange=rmse_range,
                     element='step', fill=True, alpha=0.7, linewidth=0.2, label='PSFEx')
        sns.histplot(data2['rmse'], ax=axs_rmse[row, col], color='#2c7fb8', bins=bins, binrange=rmse_range,
                    element='step', fill=True, alpha=0.7, linewidth=0.2, label='Model')  
        
        if mode == 'finetune':
            sns.histplot(data3['rmse'], ax=axs_rmse[row, col], color="#2c207b", bins=bins, binrange=rmse_range,
                        element='step', fill=True, alpha=0.7, linewidth=0.2, label='finetune')   

        axs_rmse[row, col].set_title(f'CCD {iccd}', fontsize=16)
        axs_rmse[row, col].set_xlim(rmse_range)
        axs_rmse[row, col].legend(loc='upper left', fontsize=16)
        axs_rmse[row, col].tick_params(direction='in')

    for fig, axs, name in zip(
        [fig_psnr, fig_ssim, fig_rmse],
        [axs_psnr, axs_ssim, axs_rmse],
        ['psnr', 'ssim', 'rmse']
    ):
        fig.suptitle(f'{name.upper()} Histograms Comparison', y=1.0, fontsize=20)
        fig.tight_layout()
        fig.delaxes(axs_psnr[0, 4]) if name == 'psnr' else None
        fig.delaxes(axs_psnr[3, 0]) if name == 'psnr' else None
        fig.delaxes(axs_ssim[0, 4]) if name == 'ssim' else None
        fig.delaxes(axs_ssim[3, 0]) if name == 'ssim' else None
        fig.delaxes(axs_rmse[0, 4]) if name == 'rmse' else None
        fig.delaxes(axs_rmse[3, 0]) if name == 'rmse' else None

        fig.savefig(os.path.join(save_dir, f'combined_{name}_histograms_4x5.png'),
                    dpi=300, bbox_inches='tight')
    plt.close(fig)

        
def shape_measure(stamp: np.ndarray, upsampling=False, oversampling=1) -> tuple:
    if upsampling:
        from scipy import ndimage   
        stamp = ndimage.zoom(stamp, (2, 2), order=3)
        oversampling *= 2
    galimage = galsim.Image(stamp, scale=1)
    AM = galsim.hsm.FindAdaptiveMom(galimage, strict=False)

    if AM.moments_sigma <= 0:
        return np.nan, np.nan, np.nan, np.nan, np.nan

    AM_shape = AM.observed_shape
    fwhm_am = AM.moments_sigma * 2 * np.sqrt(2 * np.log(2))
    e = AM_shape.e 
    e1 = AM_shape.e1
    e2 = AM_shape.e2
    pa = AM_shape.beta.deg % 180  # in degrees
    return (
        fwhm_am * 0.074 / oversampling,
        e,
        e1,
        e2,
        pa,
    )

def meas_iccd(data_dir, iccd):
    """
    Process a single ICCD file and return a table with measurements.
    Preserves all objects even if measurements are NaN.
    
    Parameters:
    -----------
    data_dir : str
        Directory containing the FITS files
    iccd : int
        ICCD number to process
        
    Returns:
    --------
    astropy.Table
        Table containing measurements for all objects in this ICCD
    dict
        Dictionary with NaN counts per column for this ICCD
    """
    filename = os.path.join(data_dir, f"psf_ccd{iccd}.fits")
    # print(filename)
    with fits.open(filename) as hdu:
        imgs = hdu[1].data
        tab = Table.read(filename, hdu=2)
    
    # Initialize output columns
    n_objs = len(tab)
    result = Table()
    result['iccd'] = np.full(n_objs, iccd, dtype=int)
    result['obj_id'] = tab['obj_id'] if 'obj_id' in tab.columns else np.arange(n_objs)
    result['x_pos'] = tab['X_IMAGE'] if 'X_IMAGE' in tab.columns else np.zeros(n_objs)
    result['y_pos'] = tab['Y_IMAGE'] if 'Y_IMAGE' in tab.columns else np.zeros(n_objs)
    
    # Measurement columns initialized with NaN
    float_cols = ['FWHM', 'e', 'e1', 'e2', 'pa', 
                 'FWHM_4x', 'e_4x', 'e1_4x', 'e2_4x', 'pa_4x']
    for col in float_cols:
        result[col] = np.full(n_objs, np.nan, dtype=float)
    
    # Dictionary to track NaN counts
    nan_counts = {col: 0 for col in float_cols}
    nan_objects = {col: [] for col in float_cols}
    
    # Process each image
    for i in range(n_objs):
        img = imgs[i, :, :]
        current_obj = result['obj_id'][i]
        
        try:
            # Regular measurements
            fwhm, e, e1, e2, pa = shape_measure(img, upsampling=False)
            result['FWHM'][i] = fwhm
            result['e'][i] = e
            result['e1'][i] = e1
            result['e2'][i] = e2
            result['pa'][i] = pa
            
            # 4x upsampled measurements
            fwhm_4x, e_4x, e1_4x, e2_4x, pa_4x = shape_measure(img, upsampling=True, oversampling=2)
            result['FWHM_4x'][i] = fwhm_4x
            result['e_4x'][i] = e_4x
            result['e1_4x'][i] = e1_4x
            result['e2_4x'][i] = e2_4x
            result['pa_4x'][i] = pa_4x
        except Exception as e:
            # If measurement fails, values remain NaN
            pass
        
        # Update NaN counts for this object
        for col in float_cols:
            if np.isnan(result[col][i]):
                nan_counts[col] += 1
                nan_objects[col].append(current_obj)
    
    return result, nan_counts, nan_objects

def meas_ccds(data_dir, iccd_list, output_filename):
    """
    Process a list of ICCDs and save results to a single sorted FITS table.
    Includes NaN statistics in a separate extension.
    
    Parameters:
    -----------
    data_dir : str
        Directory containing the FITS files
    iccd_list : list of int
        List of ICCD numbers to process
    output_filename : str
        Path to save the output FITS file
    """
    # Process all ICCDs
    tables = []
    all_nan_counts = []
    all_nan_objects = []
    
    for iccd in tqdm(iccd_list, desc="Processing ICCDs"):
        try:
            table, nan_counts, nan_objects = meas_iccd(data_dir, iccd)
            tables.append(table)
            all_nan_counts.append(nan_counts)
            all_nan_objects.append(nan_objects)
        except FileNotFoundError:
            print(f"Warning: File for ICCD {iccd} not found in {data_dir}")
            continue
        except Exception as e:
            print(f"Error processing ICCD {iccd}: {str(e)}")
            continue
    
    if not tables:
        raise ValueError("No valid ICCD files were processed")
    
    # Combine all tables and sort by iccd
    combined_table = vstack(tables)
    combined_table.sort(['iccd', 'obj_id'])
    
    # Create NaN statistics table
    nan_stats = Table()
    nan_stats['column'] = list(all_nan_counts[0].keys())
    
    # Calculate total NaN counts across all ICCDs
    total_nan_counts = {col: 0 for col in all_nan_counts[0]}
    for counts in all_nan_counts:
        for col, cnt in counts.items():
            total_nan_counts[col] += cnt
    nan_stats['nan_count'] = [total_nan_counts[col] for col in nan_stats['column']]
    
    # Create a string representation of affected object IDs for each column
    nan_objects_combined = {col: [] for col in all_nan_objects[0]}
    for nobjs in all_nan_objects:
        for col, objs in nobjs.items():
            nan_objects_combined[col].extend(objs)
    
    nan_stats['affected_objects'] = [
        str(objs[:100]) + ('...' if len(objs) > 100 else '') 
        for objs in nan_objects_combined.values()
    ]
    nan_stats['total_objects'] = len(combined_table)
    nan_stats['nan_fraction'] = nan_stats['nan_count'] / nan_stats['total_objects']
    
    # Save to FITS with two extensions
    hdu_primary = fits.PrimaryHDU()
    hdu_measurements = fits.table_to_hdu(combined_table)
    hdu_nan_stats = fits.table_to_hdu(nan_stats)
    
    hdul = fits.HDUList([
        hdu_primary,
        hdu_measurements,
        hdu_nan_stats
    ])
    
    # Set extension names
    hdu_measurements.name = 'MEASUREMENTS'
    hdu_nan_stats.name = 'NAN_STATS'
    
    hdul.writeto(output_filename, overwrite=True)
    
    print(f"Saved results for {len(tables)} ICCDs to {output_filename}")
    # print("\nNaN Statistics Summary:")
    # nan_stats.pprint(max_width=-1)
    
    return combined_table


from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import griddata
from astropy.table import Table, join
def residual_map(tab_t_path, tab_r_path, col_name, output_prefix,
                    method_label="Model", 
                    grid_shape=(30, 30), 
                    extent=(0, 9216, 0, 9232),   
                    interp="nearest",
                    res_vlim=0.02,
                    save_residual_fits = True,
                    save_perccd_rms_csv=True
                    ):
    im_dir = os.path.join(output_prefix, 'plot_imshow'); os.makedirs(im_dir, exist_ok=True)

    tab_true = Table.read(tab_t_path, hdu=1) 
    tab_recon = Table.read(tab_r_path, hdu=1)
    
    keys = ['iccd', 'x_pos', 'y_pos']
    for k in keys:
        if k not in tab_true.colnames or k not in tab_recon.colnames:
            raise ValueError(f"Missing join key '{k}' in tables.")
    tab = join(tab_true, tab_recon, keys=keys, table_names=('true','rec'), uniq_col_name='{col_name}_{table_name}', join_type='inner')

    # 取真值与重建列
    v_true = tab[f'{col_name}_true']
    v_rec  = tab[f'{col_name}_rec']
        
    # 规则网格（用于可视化）
    nx, ny = grid_shape
    xmin, xmax, ymin, ymax = extent
    grid_x, grid_y = np.mgrid[xmin:xmax:complex(nx), ymin:ymax:complex(ny)]
    
    # --- 1) TRUE 图 ---
    fig_true, axs_true = plt.subplots(4, 5, figsize=(24, 18))
    for ax in axs_true.ravel():
        ax.set_xticks([]); ax.set_yticks([])
        
    fig_true.subplots_adjust(left=0.02, right=0.9, bottom=0.02, top=0.98, wspace=0.0, hspace=0.0)

    # --- 2) RECON 图 ---
    fig_recon, axs_recon = plt.subplots(4, 5, figsize=(24, 18))
    for ax in axs_recon.ravel():
        ax.set_xticks([]); ax.set_yticks([])
    fig_recon.subplots_adjust(left=0.02, right=0.9, bottom=0.02, top=0.98, wspace=0.0, hspace=0.0)

    # --- 3) RESIDUAL 图 ---
    fig_res, axs_res = plt.subplots(4, 5, figsize=(24, 18))
    for ax in axs_res.ravel():
        ax.set_xticks([]); ax.set_yticks([])
    fig_res.subplots_adjust(left=0.02, right=0.9, bottom=0.02, top=0.98, wspace=0.0, hspace=0.0)
        
    # 色标范围
    vmin_res, vmax_res = -res_vlim, +res_vlim
    true_min, true_max   = np.nanmin(v_true),  np.nanmax(v_true)
    recon_min, recon_max = np.nanmin(v_rec),   np.nanmax(v_rec)
    
    # 统计、FITS 容器与缓存
    per_ccd_rms = []
    grid_res_by_ccd = {}   # 缓存每个 CCD 的残差网格（imshow 的朝向）
    pts_by_ccd      = {}   # 缓存每个 CCD 的星点 (x,y,resid) + RMS
    
    if save_residual_fits:
        hdus = [fits.PrimaryHDU()]
        hdr = hdus[0].header
        hdr['METHOD']  = method_label
        hdr['COLNAME'] = col_name
        hdr['GRIDNX']  = nx
        hdr['GRIDNY']  = ny
        hdr['XMIN']    = xmin
        hdr['XMAX']    = xmax
        hdr['YMIN']    = ymin
        hdr['YMAX']    = ymax
        hdr['VMIN']    = vmin_res
        hdr['VMAX']    = vmax_res

    # 用于添加统一 colorbar
    im_true = im_recon = im_res = None
    
    for iccd, (ind, label) in iccd_positions.items():
        row, col = ind // 5, ind % 5
        sub = tab[tab['iccd'] == iccd]  
        
        xs, ys = sub['x_pos'], sub['y_pos']
        vt = sub[f'{col_name}_true']
        vr = sub[f'{col_name}_rec']

        if 'FWHM' in col_name:
            ds = (vr**2 - vt**2) / (vt**2 + 1e-8)
            ds -= np.median(ds)
        else:
            ds = vr - vt
            
        # 星点残差 RMS（无权）
        rms = float(np.sqrt(np.nanmean(ds**2)))
        per_ccd_rms.append((iccd, rms))
        
        # 插值
        grid_true = griddata((xs, ys), vt, (grid_x, grid_y), method=interp)
        grid_recon = griddata((xs, ys), vr, (grid_x, grid_y), method=interp)
        grid_res = griddata((xs, ys), ds, (grid_x, grid_y), method=interp)
        
        # 缓存本 CCD 的残差网格与星点（注意 T：与 imshow 一致的朝向）
        grid_res_by_ccd[iccd] = grid_res.T.astype('f4')
        pts_by_ccd[iccd] = (
            np.asarray(xs, dtype=np.float32),
            np.asarray(ys, dtype=np.float32),
            np.asarray(ds, dtype=np.float32),
            rms  # 顺便把 RMS 放一起
        )
        
        # TRUE
        im_true = axs_true[row, col].imshow(
            grid_true.T, origin='lower', extent=extent, cmap='viridis', 
            vmin=true_min, vmax=true_max,
            aspect='equal'
        )

        # RECON
        im_recon = axs_recon[row, col].imshow(
            grid_recon.T, origin='lower', extent=extent, cmap='viridis', 
             vmin=recon_min, vmax=recon_max,
             aspect='equal'
        )
        
        # RESID            
        im_res = axs_res[row, col].imshow(
            grid_res.T, origin='lower', extent=extent, cmap='coolwarm',
            vmin=vmin_res, vmax=vmax_res, aspect='equal'
        )
        axs_res[row, col].text(0.05, 0.05, f'RMS={rms:.5e}',
                               transform=axs_res[row, col].transAxes,
                               ha='left', va='top',
                               bbox=dict(boxstyle='round,pad=0.2', fc='white', ec='0.7', alpha=0.8),
                               fontsize=20,
                               fontweight='bold' 
                               )      
    # 统一 colorbar（放右侧细条）
    def add_cbar(fig, im, label):
        cax = fig.add_axes([0.92, 0.10, 0.015, 0.80])
        cb = fig.colorbar(im, cax=cax)
        cb.set_label(label, fontsize=20, fontweight='bold', labelpad=15)
        cb.ax.tick_params(labelsize=20, width=1.5, length=5)
    
    add_cbar(fig_true, im_true,  fr'{col_name} (true)')
    add_cbar(fig_recon, im_recon, fr'{col_name} (recon)')
    add_cbar(fig_res, im_res, 
             fr'$\delta R^2/R^2$' if 'FWHM' in col_name else fr'$\delta {col_name}$')        
    
    # 删空白轴（与原布局一致）
    def safe_delaxes(fig, ax, r, c):
        try:
            fig.delaxes(ax[r, c])
        except Exception:
            pass
        
    for f, ax in [(fig_true, axs_true), (fig_recon, axs_recon), (fig_res, axs_res)]:
        safe_delaxes(f, ax, 0, 4)
        safe_delaxes(f, ax, 3, 0)
        
    # 保存图片（无缝+无标题）
    fig_true.savefig(os.path.join(im_dir, f"true_{col_name}.png"),  dpi=300, bbox_inches='tight')
    fig_recon.savefig(os.path.join(im_dir, f"recon_{col_name}.png"), dpi=300, bbox_inches='tight')
    fig_res.savefig(os.path.join(im_dir, f"residual_{col_name}.png"), dpi=300, bbox_inches='tight')
    plt.close(fig_true); plt.close(fig_recon); plt.close(fig_res)
                    
    # 残差 FITS（每 CCD 一层 + 星点表）
    fits_path = None
    if save_residual_fits:
        hdul = fits.HDUList(hdus)  # 含 PrimaryHDU
        for iccd, (ind, label) in sorted(iccd_positions.items(), key=lambda x: x[1][0]):
            if iccd not in grid_res_by_ccd:
                continue
            img = grid_res_by_ccd[iccd]
            xarr, yarr, resarr, rms_ccd = pts_by_ccd[iccd]

            ihdu = fits.ImageHDU(data=img, name=f'CCD{iccd}')
            ihdu.header['ICCD']   = iccd
            ihdu.header['LABEL']  = label
            ihdu.header['METHOD'] = method_label
            ihdu.header['RMS']    = rms_ccd
            ihdu.header['VMIN']   = vmin_res
            ihdu.header['VMAX']   = vmax_res
            ihdu.header['NX']     = nx
            ihdu.header['NY']     = ny
            ihdu.header['XMIN']   = xmin
            ihdu.header['XMAX']   = xmax
            ihdu.header['YMIN']   = ymin
            ihdu.header['YMAX']   = ymax
            hdul.append(ihdu)

            cols = [
                fits.Column(name='x_pos', format='E', array=xarr),
                fits.Column(name='y_pos', format='E', array=yarr),
                fits.Column(name='resid', format='E', array=resarr),
            ]
            thdu = fits.BinTableHDU.from_columns(cols, name=f'PTS{iccd}')
            thdu.header['ICCD']   = iccd
            thdu.header['LABEL']  = label
            thdu.header['METHOD'] = method_label
            hdul.append(thdu)

        fits_path = os.path.join(im_dir, f"residual_{col_name}_{method_label}.fits")
        hdul.writeto(fits_path, overwrite=True)
        
    # 保存 per-CCD RMS CSV（用于条形图/统计）
    csv_path = None
    if save_perccd_rms_csv:
        import pandas as pd
        df = pd.DataFrame(per_ccd_rms, columns=['iccd', f'RMS_{col_name}'])
        csv_path = os.path.join(im_dir, f"rms_{col_name}_{method_label}.csv")
        df.to_csv(csv_path, index=False)

    return 

def plot_histograms_fwhm_e(tab_t_path, tab_r_path, output_prefix, columns=None):
    # Read tables
    tab_true = Table.read(tab_t_path, hdu=1) 
    tab_recon = Table.read(tab_r_path, hdu=1)
    
    # Default columns to plot if not specified
    if columns is None:
        columns = ['FWHM', 'e', 'e1', 'e2', 'FWHM_4x', 'e_4x', 'e1_4x', 'e2_4x']
    
    for col_name in columns:
        # Calculate delta values
        if 'FWHM' in col_name:
            delta = (tab_recon[col_name]**2 - tab_true[col_name]**2) / (tab_true[col_name]**2 + 1e-6)
            delta_label = f'δR2/R2 ({col_name})'
        else:
            delta = tab_recon[col_name] - tab_true[col_name]
            delta_label = f'δ{col_name}'
        
        # Create figures
        fig_true, axs_true = plt.subplots(4, 5, figsize=(24, 18))
        fig_recon, axs_recon = plt.subplots(4, 5, figsize=(24, 18))
        fig_res, axs_res = plt.subplots(4, 5, figsize=(24, 18))
        
        # Plot histograms for each ICCD
        for iccd, (row, col) in iccd_positions.items():
            pos = get_subplot_pos(iccd)
            row = (pos - 1) // 5
            col = (pos - 1) % 5
            
            # Get data for current ICCD
            mask_true = tab_true['iccd'] == iccd
            mask_recon = tab_recon['iccd'] == iccd
            mask = mask_true & mask_recon
            
            t_data = np.array(tab_true[col_name][mask], dtype=np.float64)
            r_data = np.array(tab_recon[col_name][mask], dtype=np.float64)
            d_data = np.array(delta[mask], dtype=np.float64)
            
            # Plot true data
            t_median = np.nanmedian(t_data)
            t_std = np.nanstd(t_data)
            
            sns.histplot(t_data, ax=axs_true[row, col], kde=True, color='skyblue', bins=20)
            axs_true[row, col].axvline(t_median, color='r', linestyle='--', linewidth=1.5)
            axs_true[row, col].set_title(f'CCD {iccd}\nMedian: {t_median:.4f} ± {t_std:.4f}')
            axs_true[row, col].tick_params(direction='in')
            
            # Plot recon data
            r_median = np.nanmedian(r_data)
            r_std = np.nanstd(r_data)
            
            sns.histplot(r_data, ax=axs_recon[row, col], kde=True, color='skyblue', bins=20)
            axs_recon[row, col].axvline(r_median, color='r', linestyle='--', linewidth=1.5)
            axs_recon[row, col].set_title(f'CCD {iccd}\nMedian: {r_median:.4f} ± {r_std:.4f}')
            axs_recon[row, col].tick_params(direction='in')
            
            # Plot residual/delta data 
            d_median = np.nanmedian(d_data)
            d_std = np.nanstd(d_data)
            
            sns.histplot(d_data, ax=axs_res[row, col], kde=True, color='skyblue', bins=20)
            axs_res[row, col].axvline(d_median, color='r', linestyle='--', linewidth=1.5)
            axs_res[row, col].set_title(f'CCD {iccd}\nMedian: {d_median:.4f} ± {d_std:.4f}')
            axs_res[row, col].tick_params(direction='in')

        hist_dir = os.path.join(output_prefix, 'histograms')
        os.makedirs(hist_dir, exist_ok=True)
                
        # Adjust and save true histograms
        fig_true.suptitle(f'True PSF for {col_name}', fontsize=20)
        fig_true.delaxes(axs_true[0, 4])
        fig_true.delaxes(axs_true[3, 0])
        fig_true.subplots_adjust(left=0.05, right=0.9, top=0.9, bottom=0.05)
        fig_true.savefig(f"{hist_dir}/true_histograms_{col_name}.png", dpi=300, bbox_inches='tight')
        
        # Adjust and save recon histograms
        fig_recon.suptitle(f'Recon PSF for {col_name}', fontsize=20)
        fig_recon.delaxes(axs_recon[0, 4])
        fig_recon.delaxes(axs_recon[3, 0])
        fig_recon.subplots_adjust(left=0.05, right=0.9, top=0.9, bottom=0.05)
        fig_recon.savefig(f"{hist_dir}/recon_histograms_{col_name}.png", dpi=300, bbox_inches='tight')
        
        # Adjust and save residual histograms
        fig_res.suptitle(f'{delta_label}', fontsize=20)
        fig_res.delaxes(axs_res[0, 4])
        fig_res.delaxes(axs_res[3, 0])
        fig_res.subplots_adjust(left=0.05, right=0.9, top=0.9, bottom=0.05)
        fig_res.savefig(f"{hist_dir}/residual_histograms_{col_name}.png", dpi=300, bbox_inches='tight')
        
        # Close figures to free memory
        plt.close(fig_true)
        plt.close(fig_recon)
        plt.close(fig_res)

import numpy as np
from astropy.table import Table
from astropy.stats import sigma_clip
import os
import warnings
import matplotlib.colors as mcolors
import matplotlib.ticker as mticker

def save_delta_stats(tab_t_path, tab_r_path, output_dir, params=None):
    """增强版PSF参数差异统计函数
    
    Args:
        tab_t_path: 真实PSF表格路径
        tab_r_path: 重建PSF表格路径  
        output_dir: 输出目录
        params: 要处理的参数列表(默认为None时处理全部参数)
    """
    # 抑制numpy的特定警告
    warnings.filterwarnings('ignore', 
                         message='Warning: \'partition\' will ignore the \'mask\' of the MaskedColumn')
    
    # 参数默认设置
    if params is None:
        params = ['FWHM', 'e', 'e1', 'e2', 'FWHM_4x', 'e_4x', 'e1_4x', 'e2_4x']
    
    # 读取数据
    tab_true = Table.read(tab_t_path, hdu=1)
    tab_recon = Table.read(tab_r_path, hdu=1)
    
    # 获取所有ICCD编号
    all_iccds = np.unique(np.concatenate([tab_true['iccd'], tab_recon['iccd']]))
    print(f"处理中 | 检测到ICCD传感器: {all_iccds.tolist()}")
    
    # 处理每个参数
    for param in params:
        print(f"\n正在处理参数: {param}", end=' | ')
        
        # 计算差异值
        if 'FWHM' in param:
            delta = (tab_recon[param]**2 - tab_true[param]**2) / (tab_true[param]**2 + 1e-8)
            # delta = delta - np.median(delta)
            delta_type = "δR²/R²"
        else:
            delta = tab_recon[param] - tab_true[param]
            delta_type = f"δ{param}"
        
        results = []
        all_valid = []
        
        # 处理每个ICCD
        for iccd in all_iccds:
            mask = (tab_true['iccd'] == iccd) & (tab_recon['iccd'] == iccd)
            iccd_data = np.asarray(delta[mask], dtype=np.float64)
            
            # 数据清洗流程
            clean_data = iccd_data[~np.isnan(iccd_data)]  # 去NaN
            if len(clean_data) == 0:
                continue
                
            clipped = sigma_clip(clean_data, sigma=3)  # 3σ剔除
            valid_data = clean_data[~clipped.mask]      # 有效数据
            
            if len(valid_data) < 3:  # 至少需要3个点才统计
                continue
            
            # 获取传感器信息
            filter_info = iccd_positions.get(iccd, (None, ""))[1]
            base_filter = filter_info.rstrip('1234')
            cenwave = filter_CenWave.get(base_filter, -1)
            
            # 保存结果
            results.append({
                'ICCD': iccd,
                'Filter': filter_info,
                'CenWave': cenwave,
                'Median': np.median(valid_data),
                'Std': np.std(valid_data),
                'N': len(valid_data)
            })
            all_valid.extend(valid_data)
        
        # 生成报告
        os.makedirs(output_dir, exist_ok=True)
        generate_report(param, delta_type, results, all_valid, output_dir)
        print(f"已完成 → {os.path.join(output_dir, f'psf_δ{param.lower()}_stat.txt')}")

def generate_report(param, delta_type, results, all_data, output_dir):
    """生成格式化的统计报告"""
    # 总体统计
    overall_stats = {
        'Median': np.median(all_data) if len(all_data) > 0 else np.nan,
        'Std': np.std(all_data) if len(all_data) > 0 else np.nan,
        'N': len(all_data)
    }
    
    # 构建输出内容
    header = ("{:<6} {:<6} {:<8} {:>12} {:>12} {:>8}"
              .format("ICCD", "Filter", "CenWave", "Median", "StdDev", "N"))
    sep = '-' * 60
    lines = [
        f"{delta_type} Analysis for {param}",
        '=' * 60,
        header,
        sep
    ]
    
    # 添加每个ICCD的数据
    for res in sorted(results, key=lambda x: x['ICCD']):
        lines.append("{ICCD:<6d} {Filter:<6} {CenWave:<8d} {Median:>12.6f} {Std:>12.6f} {N:>8d}"
                    .format(**res))
    
    # 添加总体统计
    lines.extend([
        sep,
        "{:<6} {:<6} {:<8} {:>12.6f} {:>12.6f} {:>8d}"
        .format("ALL", "-", "-", overall_stats['Median'], overall_stats['Std'], overall_stats['N']),
        '=' * 60
    ])
    
    # 写入文件
    output_file = os.path.join(output_dir, f'psf_δ{param.lower()}_stat.txt')
    with open(output_file, 'w') as f:
        f.write('\n'.join(lines))

def plot_compare_psf(psf_dir, meas_fits, title, coord_file, output_dir, compare_dir=None):
    fig, axs = plt.subplots(4, 5, figsize=(25, 15))
    if compare_dir:
        fig_diff, axs_diff = plt.subplots(4, 5, figsize=(25, 15))
    
    os.makedirs(output_dir, exist_ok=True)
    
    coord_data = np.genfromtxt(coord_file, dtype=None, names=True, encoding='utf-8')
    coord_dict = {row['ICCD']: (row['OBJ_ID'], row['X_IMAGE'], row['Y_IMAGE']) 
                 for row in coord_data}
    
    meas_tab = Table.read(meas_fits, hdu=1)

    # ============================================
    # 预先计算统一的color scale
    # ============================================
    # 预先加载所有参考PSF
    ref_imgs_dict = {}
    if compare_dir:
        for iccd in iccd_positions.keys():
            psf_path = os.path.join(compare_dir, f'psf_ccd{iccd}.fits')
            if os.path.exists(psf_path):
                with fits.open(psf_path) as hdu:
                    ref_imgs_dict[iccd] = hdu[1].data
            else:
                print(f"警告: 参考PSF文件不存在: {psf_path}")
                ref_imgs_dict[iccd] = None
                
    all_diffs = []            
    if compare_dir:
        for iccd, (ind, label) in iccd_positions.items():
            obj_id, x_pos, y_pos = coord_dict[iccd]
            psf_path = os.path.join(psf_dir, f'psf_ccd{iccd}.fits')
            
            with fits.open(psf_path) as hdu:
                imgs = hdu[1].data
                t_tab = meas_tab[(meas_tab['iccd'] == iccd) & (meas_tab['obj_id'] == obj_id)]
                matched_idx = t_tab['obj_id'] == obj_id
                img_idx = np.where(matched_idx)[0][0]
                img = imgs[img_idx]
                
                if ref_imgs_dict[iccd] is not None:
                    ref_img = ref_imgs_dict[iccd][img_idx]
                    diff_img = img.astype(np.float32) - ref_img.astype(np.float32)
                    all_diffs.append(diff_img)
        
        # 计算全局统一的color scale
        if all_diffs:
            global_max_diff = np.max([np.max(np.abs(d)) for d in all_diffs])
            vmin_diff, vmax_diff = -global_max_diff, global_max_diff
            print(f"统一残差color scale: [{vmin_diff:.2e}, {vmax_diff:.2e}]")
        else:
            vmin_diff, vmax_diff = -1e-4, 1e-4
        
    im_diff_sample = None
    for iccd, (ind, label) in iccd_positions.items():
        row = ind // 5
        col = ind % 5
        
        obj_id, x_pos, y_pos = coord_dict[iccd]
        psf_path = os.path.join(psf_dir, f'psf_ccd{iccd}.fits')
        
        # 获取当前CCD的波段和中心波长信息
        band = label[:-1]  
        cenwave = filter_CenWave.get(band, 0)
        
        with fits.open(psf_path) as hdu:
            imgs = hdu[1].data
            t_tab = meas_tab[(meas_tab['iccd'] == iccd) & (meas_tab['obj_id'] == obj_id)]         
                   
            fwhm_4x = t_tab['FWHM_4x'][0]
            e = t_tab['e_4x'][0]
                
            matched_idx = t_tab['obj_id'] == obj_id
            img_idx = np.where(matched_idx)[0][0]
            img = imgs[img_idx]

            # Plot the PSF image
            norm = simple_norm(img, 'log', percent=100)
            im = axs[row, col].imshow(img, 
                                      norm=norm, 
                                      cmap="viridis", 
                                      origin='lower', 
                                      aspect="equal")
            
            # 使用LogNorm确保colorbar正确显示对数刻度
            # vmin_psf = max(img.min(), 1e-6)  # 避免log(0)
            # vmax_psf = img.max()

            # im = axs[row, col].imshow(
            #     img, 
            #     norm=LogNorm(vmin=vmin_psf, vmax=vmax_psf),
            #     cmap="viridis", 
            #     origin='lower', 
            #     aspect="equal"
            # )
            axs[row, col].set_title(f'CCD{iccd} {label} {cenwave}Å\nFWHM={fwhm_4x:.3f}″({fwhm_4x/0.074:.2f} pix), e={e:.3f}', 
                                   fontsize=10)
            axs[row, col].axis('off')
            
            cbar = plt.colorbar(im, ax=axs[row, col], fraction=0.046, pad=0.04) 
            
            tick_values = [1e-5, 1e-4, 1e-3, 1e-2]
            # tick_values = [t for t in tick_values if img.min() <= t <= img.max()+1e-10]
            cbar.set_ticks(tick_values)
            cbar.set_ticklabels([f'$10^{{{int(np.log10(t))}}}$' for t in tick_values])
            cbar.ax.tick_params(labelsize=8)
                        
            # 如果提供了比较目录，绘制差值图
            if compare_dir and ref_imgs_dict[iccd] is not None:
                ref_img = ref_imgs_dict[iccd][img_idx]
                
                # 计算差值（转换为浮点以避免溢出）
                diff_img = img.astype(np.float32) - ref_img.astype(np.float32)
                
                # 计算统计信息
                max_diff = np.max(np.abs(diff_img))
                rms_diff = np.sqrt(np.mean(diff_img**2))
                
                # 绘制差值图
                # im_diff = axs_diff[row, col].imshow(diff_img, cmap='seismic', 
                #                                    vmin=-max_diff, vmax=max_diff,
                #                                    origin='lower', aspect="equal")
                
                # 使用全局统一的color scale
                im_diff = axs_diff[row, col].imshow(
                    diff_img, 
                    cmap='seismic',      # ✓ diverging colormap
                    vmin=vmin_diff,      # ✓ 统一范围（全局变量）
                    vmax=vmax_diff,      # ✓ 统一范围（全局变量）
                    origin='lower', 
                    aspect="equal"
                )
                
                if im_diff_sample is None:
                    im_diff_sample = im_diff
                    
                axs_diff[row, col].set_title(f'CCD{iccd} {label}\nRMS={rms_diff:.2e}', 
                                           fontsize=10)
                axs_diff[row, col].axis('off')
                plt.colorbar(im_diff, ax=axs_diff[row, col], fraction=0.046, pad=0.04)  
                          
    fig.subplots_adjust(left=0.05, right=0.9, top=0.9, bottom=0.05)
    fig.delaxes(axs[0, 4]) 
    fig.delaxes(axs[3, 0]) 
    fig.savefig(os.path.join(output_dir, f'{title.replace(" ", "_")}.png'), bbox_inches='tight', dpi=300)
    plt.close()


    if compare_dir:
        fig_diff.subplots_adjust(left=0.05, right=0.9, top=0.9, bottom=0.05)
        fig_diff.delaxes(axs_diff[0, 4]) 
        fig_diff.delaxes(axs_diff[3, 0]) 
        # fig_diff.text(0.5, 0.98, 
        #         f'Residuals (seismic colormap, uniform scale: ±{global_max_diff:.2e})',
        #         ha='center', va='top', fontsize=14, fontweight='bold')
        
        # # 添加统一的colorbar（右侧）
        # if im_diff_sample is not None:
        #     cbar_ax_diff = fig_diff.add_axes([0.94, 0.15, 0.02, 0.7])
        #     cbar_diff = fig_diff.colorbar(im_diff_sample, cax=cbar_ax_diff)
        #     cbar_diff.set_label('Residual', fontsize=14, rotation=270, labelpad=25)
        #     cbar_diff.ax.tick_params(labelsize=10)
            
        fig_diff.savefig(os.path.join(output_dir, f'{title.replace(" ", "_")}_Difference.png'), bbox_inches='tight', dpi=300)
        plt.close(fig_diff)
                

import pandas as pd
import matplotlib.pyplot as plt

def plot_comparison_stats(PSFEx_stat_file, model_stat_file, save_prefix, metric_name=None, mode='', finetune_stat_file=''):
    from io import StringIO
    import re
    os.makedirs(save_prefix, exist_ok=True)
    def clean_filename(s):
        # 替换所有非法字符为下划线（只保留字母数字中文）
        return re.sub(r'[\\/*?:"<>|/\s]+', '_', s)

    def read_stat_with_title(filepath):
        with open(filepath, 'r') as f:
            lines = f.readlines()

        # 提取 δR²/R² 作为标题
        raw_title = lines[0].strip()
        title = raw_title.split('Analysis')[0].strip()

        # 查找数据起始行
        data_start = next(i for i, line in enumerate(lines) if line.strip().startswith('ICCD')) + 1
        data_lines = []
        all_stats = {}

        for line in lines[data_start:]:
            if line.strip().startswith('---') or line.strip() == '':
                continue
            if line.strip().startswith('ALL'):
                parts = line.strip().split()
                all_stats['Median'] = float(parts[3])
                all_stats['StdDev'] = float(parts[4])
                continue
            if line.strip().startswith('='):  # 防止把尾部的分隔线当数据
                continue
            data_lines.append(line)

        # 手动指定列名
        column_names = ['ICCD', 'Filter', 'CenWave', 'Median', 'StdDev', 'N']
        df = pd.read_csv(StringIO(''.join(data_lines)), sep='\s+', names=column_names)

        return title, df, all_stats

    title, df1, all1 = read_stat_with_title(PSFEx_stat_file)
    _, df2, all2 = read_stat_with_title(model_stat_file)
    
    save_name = clean_filename(metric_name) if metric_name else clean_filename(title)

    # CenWave Grouped Mean
    df1_mean = df1.groupby('CenWave')[['Median', 'StdDev']].mean().reset_index()
    df2_mean = df2.groupby('CenWave')[['Median', 'StdDev']].mean().reset_index()
        

    fig, axs = plt.subplots(2, 1, figsize=(6, 8))
    axs[0].scatter(df1_mean['CenWave'], df1_mean['Median'], label='PSFEx', color='blue')
    axs[1].scatter(df1_mean['CenWave'], df1_mean['StdDev'], label='PSFEx', color='blue')
    axs[0].scatter(df2_mean['CenWave'], df2_mean['Median'], label='Model', color='orange', marker='x')
    axs[1].scatter(df2_mean['CenWave'], df2_mean['StdDev'], label='Model', color='orange', marker='x')
            
    if mode == 'finetune':
        _,df3,_ = read_stat_with_title(finetune_stat_file)
        df3_mean = df3.groupby('CenWave')[['Median', 'StdDev']].mean().reset_index()
            
        axs[0].scatter(df3_mean['CenWave'], df3_mean['Median'], label='finetune', color='green', marker='+')    
        axs[1].scatter(df3_mean['CenWave'], df3_mean['StdDev'], label='finetune', color='green', marker='+')
        

    axs[0].set_xlabel('CenWave')
    axs[0].set_ylabel('Mean Median')
    axs[0].legend()


    axs[1].set_xlabel('CenWave')
    axs[1].set_ylabel('Mean StdDev')
    axs[1].legend()

    plt.tight_layout(rect=[0, 0, 1, 0.95])

    axs[0].tick_params(direction='in')
    axs[1].tick_params(direction='in')
    print('saving:', f'{save_prefix}/med_std_cenwave_{save_name}.png')

    plt.savefig(f'{save_prefix}/med_std_cenwave_{save_name}.png', dpi=300)
    plt.close()


def plot_test(psfex_dir, model_dir):
    
    plot_compare_psnr_ssim_rmse(psfex_txt = f'{psfex_dir}/stat/psnr_ssim_rmse_full_data.txt', 
                                model_txt = f'{model_dir}/stat/psnr_ssim_rmse_full_data.txt', 
                                save_dir = model_dir)
      
    for col_name in ['FWHM', 'e', 'e1', 'e2', 'FWHM_4x', 'e_4x', 'e1_4x', 'e2_4x']:
        residual_map(model_dir + '/meas_truePSFs.fits', 
                        model_dir + '/meas_reconPSFs.fits', 
                        col_name, 
                        model_dir)

    for col_name in ['fwhm', 'e', 'e1', 'e2', 'fwhm_4x', 'e_4x', 'e1_4x', 'e2_4x']:        
        plot_comparison_stats(f'{psfex_dir}/stat/psf_δ{col_name}_stat.txt', 
                             f'{model_dir}/stat/psf_δ{col_name}_stat.txt', 
                             model_dir + '/stat',
                             metric_name=col_name)

    plot_compare_psf(model_dir + '/true_psfs', 
                     model_dir + "/meas_truePSFs.fits", 
                     'True PSFs', 
                     f'{psfex_dir}/psf_compare_imgs/objects_center_xy.txt', 
                     model_dir + '/psf_compare_imgs',
                     )
    
    plot_compare_psf(model_dir + '/recon_psfs', 
                     model_dir + "/meas_reconPSFs.fits", 
                     'Recon PSFs', 
                     f'{psfex_dir}/psf_compare_imgs/objects_center_xy.txt',
                     model_dir + '/psf_compare_imgs',
                     compare_dir = model_dir + '/true_psfs'
                     )

def plot_test_finetune(psfex_dir, model_dir, finetune_dir, mode='finetune'):
  
    for col_name in ['FWHM', 'e', 'e1', 'e2']: 
        residual_map(finetune_dir + '/meas_truePSFs.fits', 
                        finetune_dir + '/meas_reconPSFs.fits', 
                        col_name, 
                        finetune_dir)
        
    for col_name in ['fwhm', 'e', 'e1', 'e2']:        
        plot_comparison_stats(f'{psfex_dir}/stat//psf_δ{col_name}_stat.txt', 
                             f'{model_dir}/stat/psf_δ{col_name}_stat.txt', 
                             finetune_dir,
                             metric_name=col_name,
                             mode = mode,
                             finetune_stat_file = f'{finetune_dir}/stat/psf_δ{col_name}_stat.txt')

    plot_compare_psf(finetune_dir + '/true_psfs', finetune_dir + "/meas_truePSFs.fits", 'True PSFs', 
                     f'{psfex_dir}/psf_compare_imgs/objects_center_xy.txt', 
                     finetune_dir + '/psf_compare_imgs')
    plot_compare_psf(finetune_dir + '/recon_psfs', finetune_dir + "/meas_reconPSFs.fits", 'Recon PSFs', 
                     f'{psfex_dir}/psf_compare_imgs/objects_center_xy.txt',
                     finetune_dir + '/psf_compare_imgs')