quantify_ref_vs_integration.py

#!/usr/bin/env python3
"""
Comparison of [M]R (Reference Library) vs [M]I (Integration/Amplitude) Methods

This script computes metabolite concentrations using two approaches:
1. [M]R: Uses theoretical dilution ratio for Scale(M), measured Scale(TSP)
   Formula: [M]R = [M]Mref × Scale(M)_theoretical / Scale(TSP)
   
2. [M]I: Uses measured peak amplitude or integration for Scale(M), measured Scale(TSP)
   Formula: [M]I = [M]Mref × Scale(M)_measured / Scale(TSP)
   Where Scale(M)_measured = (Amplitude or Integration ratio)

The user can choose between:
- lorentzian: Use Lorentzian fitted amplitude
- integration: Use numerical integration of peak area
"""

import nmrglue as ng
import numpy as np
from scipy.ndimage import gaussian_filter1d
from scipy.signal import find_peaks
from lmfit import Model
from lmfit.models import LorentzianModel, ConstantModel
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pandas as pd
import os
import argparse


def read_and_process(filepath):
    """Read JCAMP-DX file and return ppm, magnitude spectrum"""
    dic, data = ng.jcampdx.read(filepath)
    data_real = data[0]
    data_imag = data[1] if len(data) > 1 else [0] * len(data[0])
    magnitude = np.sqrt(data_real**2 + np.array(data_imag)**2)
    smooth = gaussian_filter1d(magnitude, sigma=2)
    
    sfo1 = float(dic['$SFO1'][0])
    o1_hz = float(dic['$O1'][0])
    sw_hz = float(dic['$SWH'][0])
    o1_ppm = o1_hz / sfo1
    sw_ppm = sw_hz / sfo1
    ppm = np.linspace(o1_ppm + sw_ppm/2, o1_ppm - sw_ppm/2, len(magnitude))
    return ppm, smooth


def find_tsp_peak(ppm, intensity):
    """Find TSP reference peak position"""
    mask = (ppm >= -0.5) & (ppm <= 0.5)
    return ppm[mask][np.argmax(intensity[mask])]


def integrate_peak(ppm, intensity, region):
    """Integrate a spectral region"""
    mask = (ppm >= region[0]) & (ppm <= region[1])
    return abs(np.trapz(intensity[mask], ppm[mask]))


def fit_region(x, y, centers, n_peaks, sigma_guess=None):
    """Fit Lorentzian(s) to data"""
    if n_peaks == 1:
        model = LorentzianModel() + ConstantModel()
        params = model.make_params()
        params['amplitude'].set(value=np.max(y) * 0.01, min=0)
        params['center'].set(value=centers[0], min=centers[0]-0.05, max=centers[0]+0.05)
        params['sigma'].set(value=0.003, min=0.001, max=0.02)
        params['c'].set(value=np.min(y))
    else:
        model = None
        params = None
        for i in range(n_peaks):
            prefix = f'p{i+1}_'
            if model is None:
                model = LorentzianModel(prefix=prefix)
                params = model.make_params()
            else:
                model = model + LorentzianModel(prefix=prefix)
                params.update(model.make_params())
            
            params[f'{prefix}amplitude'].set(value=np.max(y) * 0.01 / n_peaks, min=0)
            params[f'{prefix}center'].set(value=centers[i], min=centers[i]-0.05, max=centers[i]+0.05)
            if sigma_guess:
                params[f'{prefix}sigma'].set(value=sigma_guess, min=0.001, max=0.02)
            else:
                params[f'{prefix}sigma'].set(value=0.003, min=0.001, max=0.02)
        
        model = model + ConstantModel()
        params.update(model.make_params())
        params['c'].set(value=np.min(y))
    
    result = model.fit(y, params, x=x)
    return result, model


def detect_peaks_in_region(ppm, intensity, region, expected_centers, n_peaks=1):
    """Detect peaks in spectral region and match to expected positions"""
    mask = (ppm >= region[0]) & (ppm <= region[1])
    x_region = ppm[mask]
    y_region = intensity[mask]
    
    if len(x_region) == 0 or np.max(y_region) <= 0:
        return expected_centers[:n_peaks]
    
    max_intensity = np.max(y_region)
    height_threshold = max_intensity * 0.05
    min_distance = len(x_region) // (n_peaks * 3)
    
    peaks_idx, _ = find_peaks(y_region, height=height_threshold, distance=max(10, min_distance))
    
    if len(peaks_idx) == 0:
        return expected_centers[:n_peaks]
    
    peak_positions = x_region[peaks_idx]
    peak_heights = y_region[peaks_idx]
    
    tolerance = 0.30
    detected = []
    used_peak_idx = set()
    
    for expected in expected_centers[:n_peaks]:
        distances = np.abs(peak_positions - expected)
        within_tolerance = (distances < tolerance) & ~np.isin(np.arange(len(peak_positions)), list(used_peak_idx))
        
        if np.any(within_tolerance):
            valid_indices = np.where(within_tolerance)[0]
            valid_heights = peak_heights[valid_indices]
            tallest_local_idx = np.argmax(valid_heights)
            best_idx = valid_indices[tallest_local_idx]
            detected.append(peak_positions[best_idx])
            used_peak_idx.add(best_idx)
        else:
            detected.append(expected)
    
    return detected


def quantify_metabolite_both_methods(met_name, met_info, base_dir, output_dir, method='lorentzian'):
    """
    Quantify metabolite using both [M]R and [M]I methods
    
    Parameters:
    -----------
    method : str
        'lorentzian' - use Lorentzian fitted amplitude
        'integration' - use numerical integration
    """
    folder = met_info['folder']
    regions = met_info['regions']
    region_peaks = met_info['region_peaks']
    region_protons = met_info['region_protons']
    region_names = met_info['region_names']
    files = met_info['files']
    
    folder_path = os.path.join(base_dir, folder)
    available_files = {k: v for k, v in files.items() if os.path.exists(os.path.join(folder_path, f"{k}.dx"))}
    
    if not available_files:
        return None
    
    # Use highest concentration file as reference
    ref_fileno = min(available_files.keys())
    ref_conc = available_files[ref_fileno]
    
    # Load reference
    ppm_ref, spec_ref = read_and_process(os.path.join(folder_path, f"{ref_fileno}.dx"))
    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))
    
    # Store results for each region
    region_results_MR = [[] for _ in regions]  # [M]R results
    region_results_MI = [[] for _ in regions]  # [M]I results
    
    for fileno, true_conc in sorted(available_files.items()):
        if fileno == ref_fileno:
            # Reference file
            for i in range(len(regions)):
                region_results_MR[i].append({
                    'fileno': fileno, 'true': true_conc, 'calc': ref_conc, 
                    'recovery': 100.0, 'scale_m': 1.0, 'scale_tsp': 1.0
                })
                region_results_MI[i].append({
                    'fileno': fileno, 'true': true_conc, 'calc': ref_conc,
                    'recovery': 100.0, 'scale_m': 1.0, 'scale_tsp': 1.0
                })
            continue
        
        # Load sample
        try:
            ppm_samp, spec_samp = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
        except:
            continue
        
        tsp_samp = find_tsp_peak(ppm_samp, spec_samp)
        ppm_samp_corr = ppm_samp - tsp_samp
        tsp_area_samp = integrate_peak(ppm_samp_corr, spec_samp, (-0.2, 0.2))
        scale_tsp = tsp_area_samp / tsp_area_ref
        
        # Theoretical Scale(M) for [M]R
        scale_m_theoretical = true_conc / ref_conc
        
        # Process each region
        for i, (region, peaks) in enumerate(zip(regions, region_peaks)):
            mask_ref = (ppm_ref_corr >= region[0]) & (ppm_ref_corr <= region[1])
            mask_samp = (ppm_samp_corr >= region[0]) & (ppm_samp_corr <= region[1])
            
            if not np.any(mask_ref) or not np.any(mask_samp):
                continue
            
            x_ref, y_ref = ppm_ref_corr[mask_ref], spec_ref[mask_ref] / tsp_area_ref
            x_samp, y_samp = ppm_samp_corr[mask_samp], spec_samp[mask_samp] / tsp_area_samp
            
            n_peaks = len(peaks)
            
            try:
                # Fit reference
                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
                
                # Fit sample
                detected = detect_peaks_in_region(ppm_samp_corr, spec_samp/tsp_area_samp, region, peaks, n_peaks)
                result_samp, _ = fit_region(x_samp, y_samp, detected, n_peaks, ref_sigma)
                
                if 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, n_peaks+1)])
                
                # Compute Scale(M) measured
                if method == 'lorentzian':
                    scale_m_measured = samp_amp / ref_amp if ref_amp > 1e-10 else 0
                else:  # integration
                    samp_integral = abs(np.trapz(y_samp, x_samp))
                    ref_integral = abs(np.trapz(y_ref, x_ref))
                    scale_m_measured = samp_integral / ref_integral if ref_integral > 1e-10 else 0
                
                # [M]R = [M]Mref × Scale(M)_theoretical / Scale(TSP)
                calc_MR = ref_conc * scale_m_theoretical / scale_tsp
                recovery_MR = 100 * calc_MR / true_conc if true_conc > 0 else 0
                
                # [M]I = [M]Mref × Scale(M)_measured
                # Note: scale_m_measured already includes TSP correction since
                # both amplitudes come from TSP-normalized spectra
                calc_MI = ref_conc * scale_m_measured
                recovery_MI = 100 * calc_MI / true_conc if true_conc > 0 else 0
                
                region_results_MR[i].append({
                    'fileno': fileno, 'true': true_conc, 'calc': calc_MR,
                    'recovery': recovery_MR, 'scale_m': scale_m_theoretical, 'scale_tsp': scale_tsp
                })
                region_results_MI[i].append({
                    'fileno': fileno, 'true': true_conc, 'calc': calc_MI,
                    'recovery': recovery_MI, 'scale_m': scale_m_measured, 'scale_tsp': scale_tsp
                })
                
            except Exception as e:
                print(f"    Error processing {met_name} file {fileno} region {i}: {e}")
                continue
    
    # Combine results across regions (proton-weighted)
    def combine_results(region_results):
        combined = []
        for fileno in sorted(available_files.keys()):
            true_conc = available_files[fileno]
            weighted_sum = 0
            total_protons = 0
            
            for i, region_res in enumerate(region_results):
                file_result = next((r for r in region_res if r['fileno'] == fileno), None)
                if file_result and file_result['calc'] > 0:
                    protons = region_protons[i]
                    weighted_sum += file_result['calc'] * protons
                    total_protons += protons
            
            if total_protons > 0:
                combined_calc = weighted_sum / total_protons
                combined_recovery = 100 * combined_calc / true_conc if true_conc > 0 else 0
            else:
                combined_calc = 0
                combined_recovery = 0
            
            combined.append({
                'fileno': fileno, 'true': true_conc,
                'calc': combined_calc, 'recovery': combined_recovery
            })
        return combined
    
    combined_MR = combine_results(region_results_MR)
    combined_MI = combine_results(region_results_MI)
    
    # Statistics
    def get_stats(combined):
        recoveries = [r['recovery'] for r in combined]
        true_vals = [r['true'] for r in combined]
        calc_vals = [r['calc'] for r in combined]
        
        mean_recovery = np.mean(recoveries)
        std_recovery = np.std(recoveries)
        
        if len(true_vals) >= 2:
            slope, intercept = np.polyfit(true_vals, calc_vals, 1)
            r_squared = np.corrcoef(true_vals, calc_vals)[0, 1]**2
        else:
            slope, intercept, r_squared = 0, 0, 0
        
        return mean_recovery, std_recovery, slope, intercept, r_squared
    
    stats_MR = get_stats(combined_MR)
    stats_MI = get_stats(combined_MI)
    
    return {
        'name': met_name,
        'ref_conc': ref_conc,
        'ref_fileno': ref_fileno,
        'combined_MR': combined_MR,
        'combined_MI': combined_MI,
        'stats_MR': stats_MR,
        'stats_MI': stats_MI,
        'method': method
    }


def plot_comparison(results, output_dir):
    """Plot [M]R vs [M]I comparison"""
    met_name = results['name']
    combined_MR = results['combined_MR']
    combined_MI = results['combined_MI']
    method = results['method']
    
    mean_MR, std_MR, slope_MR, intercept_MR, r2_MR = results['stats_MR']
    mean_MI, std_MI, slope_MI, intercept_MI, r2_MI = results['stats_MI']
    
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # [M]R plot
    ax = axes[0]
    true_vals_MR = [r['true'] for r in combined_MR]
    calc_vals_MR = [r['calc'] for r in combined_MR]
    ax.plot(true_vals_MR, calc_vals_MR, 'bo', markersize=10, label='[M]R data')
    ax.plot(true_vals_MR, true_vals_MR, 'k--', label='Ideal y=x')
    if r2_MR > 0:
        ax.plot(true_vals_MR, np.polyval([slope_MR, intercept_MR], true_vals_MR), 'r-', 
                label=f'Fit: y={slope_MR:.3f}x+{intercept_MR:.3f}, R²={r2_MR:.4f}')
    ax.set_xlabel('True Concentration (mM)')
    ax.set_ylabel('Calculated Concentration (mM)')
    ax.set_title(f'{met_name} - [M]R Method\nMean Recovery: {mean_MR:.1f}% ± {std_MR:.1f}%')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # [M]I plot
    ax = axes[1]
    true_vals_MI = [r['true'] for r in combined_MI]
    calc_vals_MI = [r['calc'] for r in combined_MI]
    ax.plot(true_vals_MI, calc_vals_MI, 'go', markersize=10, label=f'[M]I data ({method})')
    ax.plot(true_vals_MI, true_vals_MI, 'k--', label='Ideal y=x')
    if r2_MI > 0:
        ax.plot(true_vals_MI, np.polyval([slope_MI, intercept_MI], true_vals_MI), 'r-',
                label=f'Fit: y={slope_MI:.3f}x+{intercept_MI:.3f}, R²={r2_MI:.4f}')
    ax.set_xlabel('True Concentration (mM)')
    ax.set_ylabel('Calculated Concentration (mM)')
    ax.set_title(f'{met_name} - [M]I Method ({method})\nMean Recovery: {mean_MI:.1f}% ± {std_MI:.1f}%')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    filename = f"{met_name}_MR_vs_MI_{method}.png"
    plt.savefig(os.path.join(output_dir, filename), dpi=150)
    plt.close()
    
    print(f"  Saved: {filename}")


def get_metabolite_info():
    """Return metabolite information"""
    return {
        'Alanine': {
            'folder': 'Alanine-Reference',
            'regions': [(1.40, 1.55), (3.70, 3.85)],
            'region_peaks': [[1.48], [3.79]],
            'region_protons': [3, 1],
            'region_names': ['CH3', 'α-CH'],
            'files': {10: 40.633068, 20: 20.316534, 30: 10.158267, 40: 5.079133, 
                     50: 2.539567, 60: 1.269783, 70: 0.634892}
        },
        'Valine': {
            'folder': 'Valine-Reference',
            'regions': [(0.90, 1.15), (3.55, 3.70)],
            'region_peaks': [[0.99, 1.04], [3.62]],
            'region_protons': [6, 1],
            'region_names': ['(CH3)2', 'α-CH'],
            'files': {10: 5.021349, 20: 2.251067, 30: 1.255337, 40: 0.627669,
                     50: 0.313834, 60: 0.156917, 70: 0.078459, 80: 0.039229}
        },
        'Lactate': {
            'folder': 'Lactate-Reference',
            'regions': [(1.20, 1.45), (4.05, 4.20)],
            'region_peaks': [[1.33], [4.12]],
            'region_protons': [3, 1],
            'region_names': ['CH3', 'CH-OH'],
            'files': {10: 97.685764, 70: 48.842882, 60: 24.421441, 50: 12.210720,
                     40: 6.105360, 30: 3.052680, 20: 1.526340}
        },
        'Glucose': {
            'folder': 'Glucose-Reference',
            'regions': [(5.15, 5.35)],
            'region_peaks': [[5.24]],
            'region_protons': [1],
            'region_names': ['H-1 (anomeric)'],
            'files': {10: 103.144654, 20: 51.572327, 30: 25.786163, 40: 12.893082,
                     50: 6.446541, 60: 3.223270, 70: 1.611635}
        },
    }


def main():
    parser = argparse.ArgumentParser(description='Compare [M]R vs [M]I methods')
    parser.add_argument('--method', choices=['lorentzian', 'integration'], default='lorentzian',
                       help='Method for Scale(M) measurement (default: lorentzian)')
    parser.add_argument('--output', default='quantification_results',
                       help='Output directory for plots')
    
    args = parser.parse_args()
    
    base_dir = "raw_data/Reference_Raw_Date_JCAMP-DX"
    output_dir = args.output
    os.makedirs(output_dir, exist_ok=True)
    
    met_info_dict = get_metabolite_info()
    
    print("="*80)
    print(f"Comparing [M]R vs [M]I Methods (Scale(M) via {args.method})")
    print("="*80)
    print()
    print("[M]R: Uses theoretical dilution ratio for Scale(M)")
    print("[M]I: Uses measured amplitude/integration for Scale(M)")
    print()
    
    summary_data = []
    
    for met_name in met_info_dict:
        print(f"Processing {met_name}...")
        
        results = quantify_metabolite_both_methods(
            met_name, met_info_dict[met_name], base_dir, output_dir, args.method
        )
        
        if results:
            plot_comparison(results, output_dir)
            
            mean_MR, std_MR, _, _, r2_MR = results['stats_MR']
            mean_MI, std_MI, _, _, r2_MI = results['stats_MI']
            
            summary_data.append({
                'Metabolite': met_name,
                f'[M]R_Recovery': f"{mean_MR:.1f} ± {std_MR:.1f}",
                f'[M]R_R²': f"{r2_MR:.4f}",
                f'[M]I_Recovery': f"{mean_MI:.1f} ± {std_MI:.1f}",
                f'[M]I_R²': f"{r2_MI:.4f}",
            })
            
            print(f"  [M]R: {mean_MR:.1f}% ± {std_MR:.1f}%, R²={r2_MR:.4f}")
            print(f"  [M]I: {mean_MI:.1f}% ± {std_MI:.1f}%, R²={r2_MI:.4f}")
        else:
            print(f"  Skipped (no data)")
    
    # Print summary table
    print()
    print("="*80)
    print("SUMMARY")
    print("="*80)
    print(f"{'Metabolite':<15} {'[M]R Recovery':<20} {'[M]R R²':<12} {'[M]I Recovery':<20} {'[M]I R²':<12}")
    print("-"*80)
    for row in summary_data:
        print(f"{row['Metabolite']:<15} {row['[M]R_Recovery']:<20} {row['[M]R_R²']:<12} "
              f"{row['[M]I_Recovery']:<20} {row['[M]I_R²']:<12}")
    
    print()
    print(f"Plots saved to: {output_dir}/")
    print(f"  Filenames: {{metabolite}}_MR_vs_MI_{args.method}.png")


if __name__ == "__main__":
    main()