quantify_all_metabolites.py

#!/usr/bin/env python3
"""
Absolute Quantification of All Metabolites using TSP-normalized Lineshape Fitting
"""

import nmrglue as ng
import numpy as np
from scipy.ndimage import gaussian_filter1d
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

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 get_metabolite_info():
    """Return metabolite information: name, folder, peak region, file numbers, concentrations"""
    
    # Metabolite configurations based on typical NMR chemical shifts and Excel data
    metabolites = {
        'Alanine': {
            'folder': 'Alanine-Reference',
            'region': (1.40, 1.55),
            'peaks': [(1.48, 'd')],
            'ref_conc': 40.633068,
            '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',
            'region': (0.90, 1.15),
            'peaks': [(0.99, 'd'), (1.04, 'd')],
            'ref_conc': 5.021349,
            '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',
            'region': (1.20, 1.45),
            'peaks': [(1.33, 'd')],
            'ref_conc': 97.685764,
            'files': {10: 97.685764, 20: 1.563401, 30: 3.052680, 40: 6.105360,
                     50: 12.210720, 60: 24.421441, 70: 48.842882}
        },
        'Glucose': {
            'folder': 'Glucose-Reference',
            'region': (3.20, 3.50),
            'peaks': [(3.24, 'dd'), (3.40, 'm')],
            'ref_conc': 103.144654,
            'files': {10: 103.144654, 20: 51.572327, 30: 25.786164, 40: 12.893082,
                     50: 6.446541, 60: 3.223270, 70: 1.611635}
        },
        'Arginine': {
            'folder': 'Arginine-Reference',
            'region': (1.60, 1.85),
            'peaks': [(1.70, 'm')],
            'ref_conc': 16.696725,
            'files': {10: 16.696725, 20: 8.348363, 30: 4.174181, 40: 2.087091,
                     50: 1.043545, 60: 0.521773, 70: 0.260886}
        },
        'Glutamine': {
            'folder': 'Glutamine-Reference',
            'region': (2.40, 2.60),
            'peaks': [(2.45, 'm')],
            'ref_conc': 17.0,
            'files': {10: 17.0, 20: 8.477011, 30: 4.238505, 40: 2.119253,
                     50: 1.059626, 60: 0.529813, 70: 0.264906}
        },
        'Glutamate': {
            'folder': 'Glutamate-Reference',
            'region': (2.30, 2.60),
            'peaks': [(2.35, 'm'), (2.55, 'm')],
            'ref_conc': 10.0,
            'files': {10: 10.0, 20: 5.479266, 30: 2.739633, 40: 1.369816,
                     50: 0.684908, 60: 0.342454, 70: 0.171227}
        },
        'Aspartate': {
            'folder': 'Aspartate-Reference',
            'region': (2.60, 2.90),
            'peaks': [(2.65, 'dd'), (2.80, 'dd')],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.526296, 30: 1.263148, 40: 0.631574,
                     50: 0.315787, 60: 0.157894, 70: 0.078947}
        },
        'Asparagine': {
            'folder': 'Asparagine-Reference',
            'region': (2.80, 3.00),
            'peaks': [(2.84, 'dd')],
            'ref_conc': 7.0,
            'files': {10: 7.0, 20: 3.489252, 30: 1.744626, 40: 0.872313,
                     50: 0.436157, 60: 0.218078, 70: 0.109039}
        },
        'Isoleucine': {
            'folder': 'Isoleucine-Reference',
            'region': (0.90, 1.05),
            'peaks': [(0.94, 't'), (1.00, 'd')],
            'ref_conc': 8.0,
            'files': {10: 8.0, 20: 3.871951, 30: 1.935976, 40: 0.967988,
                     50: 0.483994, 60: 0.241997, 70: 0.120999}
        },
        'Leucine': {
            'folder': 'Leucine-Reference',
            'region': (0.90, 1.05),
            'peaks': [(0.96, 'd')],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.743902, 30: 1.371951, 40: 0.685976,
                     50: 0.342988, 60: 0.171494, 70: 0.085747}
        },
        'Methionine': {
            'folder': 'Methionine-Reference',
            'region': (2.10, 2.20),
            'peaks': [(2.15, 's')],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.598861, 30: 1.299431, 40: 0.649715,
                     50: 0.324858, 60: 0.162429, 70: 0.081214}
        },
        'Phenylalanine': {
            'folder': 'Phenylalanine-Reference',
            'region': (7.20, 7.50),
            'peaks': [(7.33, 'm')],
            'ref_conc': 5.0,
            'files': {10: 5.0, 20: 2.400125, 30: 1.200063, 40: 0.600031,
                     50: 0.300016, 60: 0.150008, 70: 0.075004}
        },
        'Tyrosine': {
            'folder': 'Tyrosine-Reference',
            'region': (6.80, 7.00),
            'peaks': [(6.90, 'd')],
            'ref_conc': 2.0,
            'files': {10: 2.0, 20: 1.087256, 30: 0.543628, 40: 0.271814,
                     50: 0.135907, 60: 0.067954, 70: 0.033977}
        }
    }
    return metabolites

def quantify_metabolite(met_name, met_info, base_dir, output_dir):
    """Quantify a single metabolite and generate plots"""
    
    folder = met_info['folder']
    region = met_info['region']
    files = met_info['files']
    ref_conc = met_info['ref_conc']
    
    folder_path = os.path.join(base_dir, folder)
    if not os.path.exists(folder_path):
        print(f"  Warning: Folder {folder_path} not found, skipping...")
        return None
    
    # Check which files exist
    available_files = {}
    for fileno, conc in files.items():
        filepath = os.path.join(folder_path, f"{fileno}.dx")
        if os.path.exists(filepath):
            available_files[fileno] = conc
    
    if len(available_files) < 2:
        print(f"  Warning: Not enough files for {met_name}, skipping...")
        return None
    
    print(f"  Processing {met_name} with {len(available_files)} files...")
    
    # Find reference file (lowest file number)
    ref_fileno = min(available_files.keys())
    ref_conc = available_files[ref_fileno]
    
    # Read and fit reference
    try:
        ppm_ref, spec_ref = read_and_process(os.path.join(folder_path, f"{ref_fileno}.dx"))
    except Exception as e:
        print(f"  Error reading reference file: {e}")
        return None
    
    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))
    
    # Fit reference
    mask_ref = (ppm_ref_corr >= region[0]) & (ppm_ref_corr <= region[1])
    x_ref = ppm_ref_corr[mask_ref]
    y_ref = spec_ref[mask_ref] / tsp_area_ref
    
    if len(x_ref) == 0:
        print(f"  Warning: No data in region {region} for {met_name}")
        return None
    
    # Use appropriate model based on number of peaks
    if len(met_info['peaks']) == 1:
        model = LorentzianModel() + ConstantModel()
        pars_ref = model.make_params()
        pars_ref['amplitude'].set(value=np.max(y_ref)*0.01, min=0)
        pars_ref['center'].set(value=met_info['peaks'][0][0], min=region[0], max=region[1])
        pars_ref['sigma'].set(value=0.005, min=0.001, max=0.02)
        pars_ref['c'].set(value=np.min(y_ref))
        
        result_ref = model.fit(y_ref, pars_ref, x=x_ref)
        ref_amplitude = result_ref.params['amplitude'].value
        ref_sigma = result_ref.params['sigma'].value
        
    else:  # Two peaks
        model = LorentzianModel(prefix='p1_') + LorentzianModel(prefix='p2_') + ConstantModel()
        pars_ref = model.make_params()
        pars_ref['p1_amplitude'].set(value=np.max(y_ref)*0.01, min=0)
        pars_ref['p1_center'].set(value=met_info['peaks'][0][0], min=region[0], max=region[1])
        pars_ref['p1_sigma'].set(value=0.005, min=0.001, max=0.02)
        pars_ref['p2_amplitude'].set(value=np.max(y_ref)*0.01, min=0)
        pars_ref['p2_center'].set(value=met_info['peaks'][1][0], min=region[0], max=region[1])
        pars_ref['p2_sigma'].set(value=0.005, min=0.001, max=0.02)
        pars_ref['c'].set(value=np.min(y_ref))
        
        result_ref = model.fit(y_ref, pars_ref, x=x_ref)
        ref_amplitude_p1 = result_ref.params['p1_amplitude'].value
        ref_amplitude_p2 = result_ref.params['p2_amplitude'].value
        ref_sigma = (result_ref.params['p1_sigma'].value + result_ref.params['p2_sigma'].value) / 2
    
    # Quantify all samples
    results = []
    for fileno, true_conc in sorted(available_files.items()):
        if fileno == ref_fileno:
            calc_conc = ref_conc
            recovery = 100.0
            scale = 1.0
            r2 = result_ref.rsquared
        else:
            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))
            
            mask_samp = (ppm_samp_corr >= region[0]) & (ppm_samp_corr <= region[1])
            x_samp = ppm_samp_corr[mask_samp]
            y_samp = spec_samp[mask_samp] / tsp_area_samp
            
            if len(x_samp) == 0:
                continue
            
            try:
                if len(met_info['peaks']) == 1:
                    pars_samp = model.make_params()
                    pars_samp['amplitude'].set(value=np.max(y_samp)*0.01, min=0)
                    pars_samp['center'].set(value=ref_amplitude, min=region[0], max=region[1])
                    pars_samp['sigma'].set(value=ref_sigma, vary=False)
                    pars_samp['c'].set(value=np.min(y_samp))
                    
                    result_samp = model.fit(y_samp, pars_samp, x=x_samp)
                    scale = result_samp.params['amplitude'].value / ref_amplitude
                    r2 = result_samp.rsquared
                else:
                    pars_samp = model.make_params()
                    pars_samp['p1_amplitude'].set(value=np.max(y_samp)*0.01, min=0)
                    pars_samp['p1_center'].set(value=met_info['peaks'][0][0], min=region[0], max=region[1])
                    pars_samp['p1_sigma'].set(value=ref_sigma, vary=False)
                    pars_samp['p2_amplitude'].set(value=np.max(y_samp)*0.01, min=0)
                    pars_samp['p2_center'].set(value=met_info['peaks'][1][0], min=region[0], max=region[1])
                    pars_samp['p2_sigma'].set(value=ref_sigma, vary=False)
                    pars_samp['c'].set(value=np.min(y_samp))
                    
                    result_samp = model.fit(y_samp, pars_samp, x=x_samp)
                    scale_p1 = result_samp.params['p1_amplitude'].value / ref_amplitude_p1
                    scale_p2 = result_samp.params['p2_amplitude'].value / ref_amplitude_p2
                    scale = (scale_p1 + scale_p2) / 2
                    r2 = result_samp.rsquared
                
                calc_conc = ref_conc * scale
                recovery = 100 * calc_conc / true_conc
            except Exception as e:
                print(f"    Error fitting file {fileno}: {e}")
                continue
        
        results.append({
            'fileno': fileno,
            'true': true_conc,
            'calc': calc_conc,
            'recovery': recovery,
            'scale': scale,
            'r2': r2
        })
    
    if len(results) < 2:
        print(f"  Warning: Not enough valid fits for {met_name}")
        return None
    
    # Calculate statistics
    recoveries = [r['recovery'] for r in results]
    mean_recovery = np.mean(recoveries)
    std_recovery = np.std(recoveries)
    
    true_vals = [r['true'] for r in results]
    calc_vals = [r['calc'] for r in results]
    
    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
    
    # Generate plots
    fig = plt.figure(figsize=(16, 12))
    gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)
    
    # Plot 1: Calibration curve
    ax1 = fig.add_subplot(gs[0, :2])
    ax1.plot(true_vals, calc_vals, 'bo', markersize=10, label='Measured')
    ax1.plot(true_vals, true_vals, 'k--', label='Ideal (y=x)')
    if r_squared > 0:
        ax1.plot(true_vals, np.polyval([slope, intercept], true_vals), 'r-', 
                 label=f'Fit: y={slope:.3f}x+{intercept:.3f}, R²={r_squared:.4f}')
    ax1.set_xlabel('True Concentration (mM)', fontsize=12)
    ax1.set_ylabel('Calculated Concentration (mM)', fontsize=12)
    ax1.set_title(f'{met_name} Quantification - Calibration Curve', fontsize=14)
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Recovery plot
    ax2 = fig.add_subplot(gs[0, 2])
    file_nums = [r['fileno'] for r in results]
    colors = ['green' if 90 < r['recovery'] < 110 else 'orange' if 80 < r['recovery'] < 120 else 'red' for r in results]
    ax2.bar(range(len(file_nums)), recoveries, color=colors)
    ax2.axhline(y=100, color='k', linestyle='--')
    ax2.axhline(y=mean_recovery, color='blue', linestyle=':', label=f'Mean: {mean_recovery:.1f}%')
    ax2.set_xticks(range(len(file_nums)))
    ax2.set_xticklabels([str(f) for f in file_nums])
    ax2.set_xlabel('File Number')
    ax2.set_ylabel('Recovery (%)')
    ax2.set_title('Recovery by Sample')
    ax2.legend()
    
    # Plot 3: All spectra
    ax3 = fig.add_subplot(gs[1, :])
    colors = plt.cm.viridis(np.linspace(0, 1, len(results)))
    for i, (r, color) in enumerate(zip(results, colors)):
        fileno = r['fileno']
        try:
            ppm, spec = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            tsp_area = integrate_peak(ppm_corr, spec, (-0.2, 0.2))
            mask = (ppm_corr >= region[0]) & (ppm_corr <= region[1])
            if np.sum(mask) > 0:
                ax3.plot(ppm_corr[mask], spec[mask]/tsp_area, color=color, 
                        label=f'{fileno} ({r["true"]:.2f} mM)', linewidth=1.5)
        except:
            pass
    ax3.set_xlabel('Chemical Shift (ppm)', fontsize=12)
    ax3.set_ylabel('TSP-normalized Intensity', fontsize=12)
    ax3.set_title(f'{met_name} Spectra (TSP-normalized)', fontsize=14)
    ax3.legend(ncol=4, fontsize=9)
    ax3.set_xlim(region[1], region[0])
    
    # Individual fits for 3 representative samples
    sample_indices = [0, len(results)//2, -1] if len(results) >= 3 else list(range(min(3, len(results))))
    for plot_idx, result_idx in enumerate(sample_indices):
        if plot_idx >= 3:
            break
        ax = fig.add_subplot(gs[2, plot_idx])
        r = results[result_idx]
        fileno = r['fileno']
        
        try:
            ppm, spec = read_and_process(os.path.join(folder_path, f"{fileno}.dx"))
            tsp = find_tsp_peak(ppm, spec)
            ppm_corr = ppm - tsp
            tsp_area = integrate_peak(ppm_corr, spec, (-0.2, 0.2))
            mask = (ppm_corr >= region[0]) & (ppm_corr <= region[1])
            x = ppm_corr[mask]
            y = spec[mask] / tsp_area
            
            if len(x) > 0:
                if len(met_info['peaks']) == 1:
                    pars = model.make_params()
                    pars['amplitude'].set(value=np.max(y)*0.01, min=0)
                    pars['center'].set(value=met_info['peaks'][0][0], min=region[0], max=region[1])
                    pars['sigma'].set(value=ref_sigma, vary=False)
                    pars['c'].set(value=np.min(y))
                else:
                    pars = model.make_params()
                    pars['p1_amplitude'].set(value=np.max(y)*0.01, min=0)
                    pars['p1_center'].set(value=met_info['peaks'][0][0], min=region[0], max=region[1])
                    pars['p1_sigma'].set(value=ref_sigma, vary=False)
                    pars['p2_amplitude'].set(value=np.max(y)*0.01, min=0)
                    pars['p2_center'].set(value=met_info['peaks'][1][0], min=region[0], max=region[1])
                    pars['p2_sigma'].set(value=ref_sigma, vary=False)
                    pars['c'].set(value=np.min(y))
                
                result = model.fit(y, pars, x=x)
                ax.plot(x, y, 'b.', markersize=3, label='Data')
                ax.plot(x, result.best_fit, 'r-', linewidth=1.5, label='Fit')
                ax.set_title(f'File {fileno}: True={r["true"]:.3f} mM\nCalc={r["calc"]:.3f} mM ({r["recovery"]:.1f}%)', fontsize=11)
                ax.set_xlabel('ppm')
                ax.set_ylabel('Intensity (norm)')
                ax.legend(fontsize=8)
                ax.set_xlim(region[1], region[0])
        except Exception as e:
            ax.text(0.5, 0.5, f'Error: {e}', ha='center', va='center', transform=ax.transAxes)
    
    output_path = os.path.join(output_dir, f'{met_name}_quantification.png')
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    # Save summary
    summary = {
        'name': met_name,
        'ref_conc': ref_conc,
        'slope': slope,
        'intercept': intercept,
        'r_squared': r_squared,
        'mean_recovery': mean_recovery,
        'std_recovery': std_recovery,
        'num_samples': len(results),
        'results': results
    }
    
    return summary

def main():
    base_dir = "raw_data/Reference_Raw_Date_JCAMP-DX"
    output_dir = "quantification_results"
    
    os.makedirs(output_dir, exist_ok=True)
    
    print("="*100)
    print("ABSOLUTE QUANTIFICATION OF ALL METABOLITES")
    print("="*100)
    print()
    
    metabolites = get_metabolite_info()
    all_summaries = []
    
    for met_name, met_info in metabolites.items():
        print(f"\nProcessing {met_name}...")
        summary = quantify_metabolite(met_name, met_info, base_dir, output_dir)
        if summary:
            all_summaries.append(summary)
            print(f"  ✓ {met_name}: R²={summary['r_squared']:.4f}, Recovery={summary['mean_recovery']:.1f}% ± {summary['std_recovery']:.1f}%")
        else:
            print(f"  ✗ {met_name}: Failed")
    
    # Print overall summary
    print()
    print("="*100)
    print("OVERALL SUMMARY")
    print("="*100)
    print(f"{'Metabolite':<20} {'Ref Conc (mM)':<15} {'R²':<10} {'Mean Rec (%)':<15} {'SD (%)':<10} {'N':<5}")
    print("-"*100)
    for s in all_summaries:
        print(f"{s['name']:<20} {s['ref_conc']:<15.2f} {s['r_squared']:<10.4f} {s['mean_recovery']:<15.1f} {s['std_recovery']:<10.1f} {s['num_samples']:<5}")
    print("="*100)
    
    # Save detailed results to CSV
    with open(os.path.join(output_dir, 'quantification_summary.csv'), 'w') as f:
        f.write("Metabolite,File,True_Conc_mM,Calc_Conc_mM,Recovery_pct,R2\n")
        for s in all_summaries:
            for r in s['results']:
                f.write(f"{s['name']},{r['fileno']},{r['true']:.6f},{r['calc']:.6f},{r['recovery']:.2f},{r['r2']:.4f}\n")
    
    print(f"\nResults saved to {output_dir}/")
    print(f"  - Individual PNG plots for each metabolite")
    print(f"  - quantification_summary.csv")

if __name__ == '__main__':
    main()