quantify_mixture_optimized.py

#!/usr/bin/env python3
"""
Optimized NMR Mixture Quantification

Uses metabolite-specific reference files and clean spectral regions
to minimize errors from peak overlap and saturation effects.

Reference file selection:
- Lactate: File 70 (best linearity when excluding File 10)
- Leucine: File 30 (CH3 region only, avoids crowded α-CH)
- Phenylalanine: File 60 (Aromatic region only, avoids crowded β-CH2)
- All others: File 20

Clean region selection:
- Leucine: Only CH3 region (0.90-1.05 ppm), skip α-CH (overlaps with Ala, Arg, Ile, Gln)
- Phenylalanine: Only Aromatic region (7.20-7.55 ppm), skip β-CH2 (overlaps with Tyr)
- Isoleucine: Only aliphatic CH3 regions, skip α-CH
- Valine: Only aliphatic CH3 regions, skip α-CH
"""

import sys
sys.path.insert(0, '.')
import os
import numpy as np
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
)

# Reference file selection per metabolite
# Use File 10 (highest concentration) as reference for all metabolites
# Based on Table S5/S6 from SI
REFERENCE_SELECTION = {
    # All metabolites use File 10 (highest concentration) as reference
    # This was verified against Table S6 values
}

# Regions to skip due to overlap (use only specified clean regions)
CLEAN_REGIONS_ONLY = {
    'Leucine': [0],           # Only CH3 region, skip α-CH
    'Phenylalanine': [0],     # Only Aromatic, skip β-CH2
    'Isoleucine': [0, 1],     # Only CH3 regions, skip α-CH
    'Valine': [0],            # Only CH3 region, skip α-CH
}

# Table S6 expected values for validation (M1 and M2)
TABLE_S6_M1 = {
    'Glucose': 62.56, 'Lactate': 61.64, 'Alanine': 16.69, 'Arginine': 10.50,
    'Glutamine': 10.46, 'Glutamate': 4.53, 'Asparagine': 2.04, 'Isoleucine': 2.04,
    'Leucine': 1.49, 'Phenylalanine': 0.97, 'Aspartate': 0.85, 
    'Methionine': 0.84, 'Valine': 0.75, 'Tyrosine': 0.76
}

TABLE_S6_M2 = {
    'Glucose': 30.99, 'Lactate': 30.11, 'Alanine': 8.22, 'Arginine': 5.34,
    'Glutamine': 5.17, 'Glutamate': 2.20, 'Asparagine': 1.01, 'Isoleucine': 1.01,
    'Leucine': 0.78, 'Phenylalanine': 0.48, 'Aspartate': 0.42, 
    'Methionine': 0.78, 'Valine': 0.33, 'Tyrosine': 0.48
}


def quantify_metabolite_optimized(met_name, met_info, mixture_ppm, mixture_spec_norm, 
                                   tsp_area_mix, reference_base_dir="raw_data/Reference_Raw_Date_JCAMP-DX"):
    """
    Quantify a single metabolite using optimized reference and regions.
    """
    # Select reference file - use File 10 (highest concentration) for all
    ref_fileno = REFERENCE_SELECTION.get(met_name, 10)
    ref_conc = met_info['files'][ref_fileno]
    
    ref_file = os.path.join(reference_base_dir, met_info['folder'], f"{ref_fileno}.dx")
    if not os.path.exists(ref_file):
        return None, f"Reference file not found: {ref_file}"
    
    # Load reference
    ppm_ref, spec_ref = read_and_process(ref_file)
    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
    
    # Determine which regions to use
    if met_name in CLEAN_REGIONS_ONLY:
        region_indices = CLEAN_REGIONS_ONLY[met_name]
    else:
        region_indices = list(range(len(met_info['regions'])))
    
    # Fit each selected region
    region_results = []
    
    for i in region_indices:
        region = met_info['regions'][i]
        peaks = met_info['region_peaks'][i]
        protons = met_info['region_protons'][i]
        region_name = met_info['region_names'][i] if i < len(met_info['region_names']) else f'R{i+1}'
        
        # Reference fit
        mask_ref = (ppm_ref_corr >= region[0]) & (ppm_ref_corr <= region[1])
        if not np.any(mask_ref):
            continue
            
        x_ref, y_ref = ppm_ref_corr[mask_ref], spec_ref_norm[mask_ref]
        n_peaks = len(peaks)
        result_ref, _ = fit_region(x_ref, y_ref, peaks, n_peaks)
        
        # Get reference amplitude
        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
        
        if ref_amp < 1e-10:
            continue
        
        # Mixture fit
        mask_mix = (mixture_ppm >= region[0]) & (mixture_ppm <= region[1])
        if not np.any(mask_mix):
            continue
            
        x_mix, y_mix = mixture_ppm[mask_mix], mixture_spec_norm[mask_mix]
        
        # Detect peaks
        detected = detect_peaks_in_region(
            mixture_ppm, mixture_spec_norm, region, peaks, n_peaks
        )
        
        # Fit sample
        result_mix, _ = fit_region(x_mix, y_mix, detected, n_peaks, ref_sigma)
        
        # Get sample amplitude
        if n_peaks == 1:
            mix_amp = result_mix.params['amplitude'].value
        else:
            mix_amp = sum([result_mix.params[f'p{j}_amplitude'].value 
                          for j in range(1, n_peaks+1)])
        
        # Calculate concentration: [M] = [M]ref × scale
        # where scale = fitted amplitude ratio of TSP-normalized spectra
        # This automatically accounts for both metabolite and TSP scaling
        scale = mix_amp / ref_amp
        calc_conc = ref_conc * scale
        
        region_results.append({
            'region': region_name,
            'calc_conc': calc_conc,
            'protons': protons,
            'scale': scale,
            'ref_amp': ref_amp,
            'mix_amp': mix_amp,
            'ref_fileno': ref_fileno,
        })
    
    # Calculate proton-weighted average
    if region_results:
        weighted_sum = sum(r['calc_conc'] * r['protons'] for r in region_results)
        total_protons = sum(r['protons'] for r in region_results)
        avg_conc = weighted_sum / total_protons
        return avg_conc, region_results
    
    return None, "No valid regions"


def quantify_mixture(mixture_file, reference_base_dir="raw_data/Reference_Raw_Date_JCAMP-DX"):
    """
    Quantify all metabolites in a mixture file using optimized parameters.
    """
    filename = os.path.basename(mixture_file)
    
    # Detect which mixture based on file number
    # Pure shift spectra mapping: 10,13,16=M1; 20,23,26=M2; 30,33,36=M3; etc.
    fileno = int(os.path.splitext(filename)[0])
    mixture_map = {
        10: 'M1', 13: 'M1', 16: 'M1',
        20: 'M2', 23: 'M2', 26: 'M2',
        30: 'M3', 33: 'M3', 36: 'M3',
        40: 'M4', 43: 'M4', 46: 'M4',
        50: 'M5', 53: 'M5', 56: 'M5',
        60: 'M6', 63: 'M6', 66: 'M6',
    }
    mixture_name = mixture_map.get(fileno, f'File{fileno}')
    
    # Select expected value table
    if mixture_name == 'M1':
        expected_table = TABLE_S6_M1
    else:
        expected_table = TABLE_S6_M2  # Default to M2 for now
    
    print("="*80)
    print(f"Optimized Quantification: {filename} ({mixture_name})")
    print("="*80)
    
    # Load mixture
    ppm, spec = read_and_process(mixture_file)
    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
    
    print(f"\nTSP correction: +{-tsp:.4f} ppm")
    print(f"TSP area: {tsp_area:.2e}\n")
    
    # Get metabolite info
    met_info_dict = get_metabolite_info_v3()
    
    results = {}
    
    for met_name, met_info in met_info_dict.items():
        print(f"\n{'='*80}")
        print(f"{met_name}")
        print(f"{'='*80}")
        
        conc, details = quantify_metabolite_optimized(
            met_name, met_info, ppm_corr, spec_norm, tsp_area, reference_base_dir
        )
        
        if conc is None:
            print(f"  Error: {details}")
            continue
        
        # Print region details
        for r in details:
            print(f"  {r['region']}: {r['calc_conc']:.2f} mM "
                  f"(ref=File {r['ref_fileno']}, scale={r['scale']:.3f})")
        
        # Calculate weighted average
        if len(details) > 1:
            weighted_sum = sum(r['calc_conc'] * r['protons'] for r in details)
            total_protons = sum(r['protons'] for r in details)
            avg_conc = weighted_sum / total_protons
            print(f"  {'='*60}")
            print(f"  Weighted average: {avg_conc:.2f} mM")
        else:
            avg_conc = conc
            
        print(f"  {'='*60}")
        
        # Compare with expected if available
        if met_name in expected_table:
            expected = expected_table[met_name]
            ratio = avg_conc / expected if expected > 0 else 0
            status = "✓" if 0.8 <= ratio <= 1.2 else "⚠"
            print(f"  Expected ({mixture_name}): {expected:.2f} mM | Ratio: {ratio:.2f} {status}")
        
        results[met_name] = {
            'concentration': avg_conc,
            'regions': details,
        }
    
    # Summary table
    print(f"\n\n{'='*80}")
    print("SUMMARY")
    print(f"{'='*80}")
    print(f"{'Metabolite':<15} {'Calculated':>12} {'Expected':>12} {'Ratio':>8} {'Status':>8}")
    print("-"*80)
    
    total_calc = 0
    total_expected = 0
    
    for met_name in sorted(results.keys()):
        calc = results[met_name]['concentration']
        total_calc += calc
        
        if met_name in expected_table:
            expected = expected_table[met_name]
            total_expected += expected
            ratio = calc / expected if expected > 0 else 0
            status = "OK" if 0.8 <= ratio <= 1.2 else "CHECK"
            print(f"{met_name:<15} {calc:>12.2f} {expected:>12.2f} {ratio:>8.2f} {status:>8}")
        else:
            print(f"{met_name:<15} {calc:>12.2f} {'N/A':>12} {'N/A':>8} {'':>8}")
    
    print("-"*80)
    print(f"{'TOTAL':<15} {total_calc:>12.2f} {total_expected:>12.2f}")
    print(f"{'='*80}")
    
    return results


def main():
    """Main function to quantify mixture files."""
    import argparse
    
    parser = argparse.ArgumentParser(description='Optimized NMR mixture quantification')
    parser.add_argument('mixture_file', nargs='?', 
                       default='raw_data/Model_Mixtures-M1-M6_JCAMP-DX/10.dx',
                       help='Path to mixture JCAMP-DX file')
    
    args = parser.parse_args()
    
    results = quantify_mixture(args.mixture_file)
    return results


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