test_smart_reference_selection.py

#!/usr/bin/env python3
"""
Smart Reference Selection for NMR Quantification

Analyses reference library to select the best reference file for each metabolite
based on:
1. TSP peak consistency (should be similar across files)
2. Spectral quality (signal-to-noise)
3. Recovery accuracy across dilution series
"""

import sys
sys.path.insert(0, '.')
import os
import numpy as np
import matplotlib.pyplot as plt
from quantify_all_metabolites_v3_ref20 import (
    read_and_process, find_tsp_peak, integrate_peak,
    detect_peaks_in_region, fit_region, get_metabolite_info_v3
)


def analyze_reference_quality(met_name, met_info, base_dir="raw_data/Reference_Raw_Date_JCAMP-DX"):
    """
    Analyze quality of each reference file for a metabolite.
    Returns quality metrics for each file.
    """
    folder = os.path.join(base_dir, met_info['folder'])
    files_dict = met_info['files']
    
    results = {}
    tsp_areas = []
    
    # First pass: collect TSP areas
    for fileno, conc in files_dict.items():
        filepath = os.path.join(folder, f"{fileno}.dx")
        if not os.path.exists(filepath):
            continue
        
        try:
            ppm, spec = read_and_process(filepath)
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            tsp_area = integrate_peak(ppm_corr, spec, (-0.2, 0.2))
            max_signal = np.max(spec)
            
            results[fileno] = {
                'concentration': conc,
                'tsp_area': tsp_area,
                'max_signal': max_signal,
                'ppm': ppm_corr,
                'spec': spec / tsp_area  # TSP-normalized
            }
            tsp_areas.append(tsp_area)
        except Exception as e:
            pass
    
    if len(tsp_areas) < 2:
        return None
    
    # Calculate median TSP area (most robust reference)
    median_tsp = np.median(tsp_areas)
    
    # Score each file
    for fileno, data in results.items():
        # TSP consistency score (closer to median = better)
        tsp_ratio = data['tsp_area'] / median_tsp
        tsp_score = 1.0 - min(abs(1 - tsp_ratio), 1.0)  # 1.0 = perfect, 0.0 = terrible
        
        # Signal-to-noise estimate
        snr = data['max_signal'] / data['tsp_area']  # Higher = better signal
        
        # Concentration appropriateness (prefer mid-range concentrations)
        # Too high = potential saturation, too low = poor SNR
        conc_ratio = data['concentration'] / max(files_dict.values())
        conc_score = 1.0 - abs(0.5 - conc_ratio) * 2  # 1.0 at 50% conc, 0 at extremes
        
        results[fileno]['tsp_ratio'] = tsp_ratio
        results[fileno]['tsp_score'] = tsp_score
        results[fileno]['snr'] = snr
        results[fileno]['conc_score'] = conc_score
        results[fileno]['quality_score'] = tsp_score * 0.4 + conc_score * 0.4 + min(snr/1000, 0.2)
    
    return results


def select_best_reference(met_name, met_info, base_dir="raw_data/Reference_Raw_Date_JCAMP-DX"):
    """
    Select the best reference file for a metabolite.
    Returns (best_fileno, quality_data)
    """
    quality_data = analyze_reference_quality(met_name, met_info, base_dir)
    
    if not quality_data:
        return None, None
    
    # Find best file by quality score
    best_fileno = max(quality_data.keys(), key=lambda f: quality_data[f]['quality_score'])
    
    return best_fileno, quality_data


def test_quantification_with_reference(met_name, met_info, ref_fileno, base_dir="raw_data/Reference_Raw_Date_JCAMP-DX"):
    """
    Test quantification accuracy using a specific reference file.
    Returns R² and mean recovery.
    """
    folder = os.path.join(base_dir, met_info['folder'])
    files_dict = met_info['files']
    regions = met_info['regions']
    region_peaks = met_info['region_peaks']
    
    # Load reference
    ref_path = os.path.join(folder, f"{ref_fileno}.dx")
    ppm_ref, spec_ref = read_and_process(ref_path)
    tsp_ref = find_tsp_peak(ppm_ref, spec_ref)
    ppm_ref_corr = ppm_ref - tsp_ref
    tsp_area_ref = integrate_peak(ppm_ref_corr, spec_ref, (-0.2, 0.2))
    spec_ref_norm = spec_ref / tsp_area_ref
    ref_conc = files_dict[ref_fileno]
    
    # Fit reference regions
    ref_fits = []
    for region, peaks in zip(regions, region_peaks):
        mask = (ppm_ref_corr >= region[0]) & (ppm_ref_corr <= region[1])
        if not np.any(mask):
            ref_fits.append(None)
            continue
        x_ref, y_ref = ppm_ref_corr[mask], spec_ref_norm[mask]
        n_peaks = len(peaks)
        result_ref, _ = fit_region(x_ref, y_ref, peaks, n_peaks)
        
        if n_peaks == 1:
            ref_amp = result_ref.params['amplitude'].value
            ref_sigma = result_ref.params['sigma'].value
        else:
            ref_amp = sum(result_ref.params[f'p{j}_amplitude'].value for j in range(1, n_peaks+1))
            ref_sigma = result_ref.params['p1_sigma'].value
        
        ref_fits.append({
            'amplitude': ref_amp,
            'sigma': ref_sigma,
            'n_peaks': n_peaks,
            'peaks': peaks
        })
    
    # Quantify other files
    calc_concs = []
    true_concs = []
    
    for fileno, true_conc in files_dict.items():
        if fileno == ref_fileno:
            calc_concs.append(true_conc)
            true_concs.append(true_conc)
            continue
        
        filepath = os.path.join(folder, f"{fileno}.dx")
        if not os.path.exists(filepath):
            continue
        
        try:
            ppm, spec = read_and_process(filepath)
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            tsp_area = integrate_peak(ppm_corr, spec, (-0.2, 0.2))
            spec_norm = spec / tsp_area
            
            # Fit each region
            region_calcs = []
            for region, ref_fit in zip(regions, ref_fits):
                if ref_fit is None:
                    continue
                
                mask = (ppm_corr >= region[0]) & (ppm_corr <= region[1])
                if not np.any(mask):
                    continue
                x_samp, y_samp = ppm_corr[mask], spec_norm[mask]
                
                detected = detect_peaks_in_region(
                    ppm_corr, spec_norm, region, ref_fit['peaks'], ref_fit['n_peaks']
                )
                result_samp, _ = fit_region(x_samp, y_samp, detected, ref_fit['n_peaks'], ref_fit['sigma'])
                
                if ref_fit['n_peaks'] == 1:
                    samp_amp = result_samp.params['amplitude'].value
                else:
                    samp_amp = sum(result_samp.params[f'p{j}_amplitude'].value for j in range(1, ref_fit['n_peaks']+1))
                
                scale = samp_amp / ref_fit['amplitude'] if ref_fit['amplitude'] > 1e-10 else 0
                calc_conc_region = ref_conc * scale
                region_calcs.append(calc_conc_region)
            
            if region_calcs:
                avg_calc = np.mean(region_calcs)
                calc_concs.append(avg_calc)
                true_concs.append(true_conc)
        except:
            pass
    
    # Calculate R² and mean recovery
    if len(calc_concs) >= 3:
        # Linear regression
        true_arr = np.array(true_concs)
        calc_arr = np.array(calc_concs)
        
        # Exclude reference point for recovery calculation
        mask = true_arr != ref_conc
        if np.any(mask):
            recoveries = calc_arr[mask] / true_arr[mask] * 100
            mean_recovery = np.mean(recoveries)
        else:
            mean_recovery = 100.0
        
        # R² calculation
        if len(true_arr) > 1:
            ss_res = np.sum((true_arr - calc_arr) ** 2)
            ss_tot = np.sum((true_arr - np.mean(true_arr)) ** 2)
            r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0
        else:
            r2 = 0
        
        return r2, mean_recovery
    
    return 0, 0


def main():
    """Analyze all metabolites and recommend best reference files"""
    
    print("="*100)
    print("SMART REFERENCE SELECTION ANALYSIS")
    print("="*100)
    
    met_info_dict = get_metabolite_info_v3()
    
    recommendations = {}
    
    print("\nAnalyzing reference file quality...\n")
    print(f"{'Metabolite':<15} {'File':>6} {'TSP_Ratio':>10} {'TSP_Score':>10} {'SNR':>10} {'Qual_Score':>12} {'R²':>8} {'Recov':>8}")
    print("-"*100)
    
    for met_name, met_info in met_info_dict.items():
        # Analyze quality
        quality_data = analyze_reference_quality(met_name, met_info)
        
        if not quality_data:
            continue
        
        # Test each file as reference
        best_file = None
        best_score = -1
        best_r2 = 0
        best_recovery = 0
        
        for fileno in sorted(quality_data.keys()):
            data = quality_data[fileno]
            
            # Test quantification accuracy
            r2, recovery = test_quantification_with_reference(met_name, met_info, fileno)
            
            # Combined score: quality + accuracy
            accuracy_score = (r2 * 0.5) + (1 - min(abs(recovery - 100) / 100, 1)) * 0.5
            total_score = data['quality_score'] * 0.5 + accuracy_score * 0.5
            
            marker = ""
            if total_score > best_score:
                best_score = total_score
                best_file = fileno
                best_r2 = r2
                best_recovery = recovery
                marker = "***"
            
            print(f"{met_name:<15} {fileno:>6} {data['tsp_ratio']:>10.2f} {data['tsp_score']:>10.2f} "
                  f"{data['snr']:>10.1f} {data['quality_score']:>12.2f} {r2:>8.3f} {recovery:>7.1f}% {marker}")
        
        recommendations[met_name] = {
            'best_file': best_file,
            'quality_data': quality_data,
            'r2': best_r2,
            'recovery': best_recovery
        }
        print()
    
    # Summary
    print("\n" + "="*100)
    print("RECOMMENDATIONS")
    print("="*100)
    print(f"{'Metabolite':<15} {'Best File':>10} {'R²':>10} {'Recovery':>12} {'TSP Ratio':>12}")
    print("-"*100)
    
    file_counts = {}
    for met_name, rec in sorted(recommendations.items()):
        best_file = rec['best_file']
        file_counts[best_file] = file_counts.get(best_file, 0) + 1
        
        tsp_ratio = rec['quality_data'][best_file]['tsp_ratio']
        print(f"{met_name:<15} {best_file:>10} {rec['r2']:>10.3f} {rec['recovery']:>11.1f}% {tsp_ratio:>11.2f}")
    
    print("\n" + "="*100)
    print("SUMMARY")
    print("="*100)
    print(f"Recommended reference files:")
    for fileno, count in sorted(file_counts.items()):
        print(f"  File {fileno}: {count} metabolites")
    
    return recommendations


if __name__ == "__main__":
    recommendations = main()