test_iterative_subtraction.py

#!/usr/bin/env python3
"""
Iterative Subtraction (CLEAN-like algorithm) for mixture quantification
"""

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

def iterative_subtraction_quantification():
    """
    Iterative approach:
    1. Find metabolite with highest correlation to residual
    2. Fit concentration
    3. Subtract contribution
    4. Repeat until convergence
    """
    
    print("="*70)
    print("Iterative Subtraction (CLEAN) - File 10")
    print("="*70)
    
    # Load references
    met_info = get_metabolite_info_v3()
    ref_base = "raw_data/Reference_Raw_Date_JCAMP-DX"
    common_ppm = np.linspace(0.5, 5.5, 1000)
    
    print("\nLoading reference spectra...")
    pure_spectra = {}
    ref_concs = {}
    
    for met_name, info in met_info.items():
        ref_file = os.path.join(ref_base, info['folder'], "20.dx")
        if not os.path.exists(ref_file):
            continue
        
        try:
            ppm, spec = read_and_process(ref_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
            spec_interp = np.interp(common_ppm, ppm_corr, spec_norm)
            
            pure_spectra[met_name] = spec_interp
            ref_concs[met_name] = info['files'][20]
            
        except Exception as e:
            print(f"  Error loading {met_name}: {e}")
    
    print(f"  Loaded {len(pure_spectra)} references")
    
    # Load mixture
    print("\nLoading mixture File 10...")
    mix_file = "raw_data/Model_Mixtures-M1-M6_JCAMP-DX/10.dx"
    ppm, spec = read_and_process(mix_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
    mixture = np.interp(common_ppm, ppm_corr, spec_norm)
    
    print(f"  TSP correction: +{-tsp:.4f} ppm")
    
    # Iterative subtraction
    print("\nRunning iterative subtraction...")
    
    residual = mixture.copy()
    results = {}
    max_iter = 20
    threshold = 0.01  # Stop when correlation below this
    
    for iteration in range(max_iter):
        # Find best matching metabolite
        best_met = None
        best_corr = 0
        best_scale = 0
        
        for met_name, pure_spec in pure_spectra.items():
            if met_name in results:  # Skip already fitted
                continue
            
            # Calculate correlation
            corr = np.corrcoef(residual, pure_spec)[0, 1]
            
            # Also calculate optimal scale (least squares)
            scale = np.dot(residual, pure_spec) / (np.dot(pure_spec, pure_spec) + 1e-10)
            
            if corr > best_corr and scale > 0:
                best_corr = corr
                best_met = met_name
                best_scale = scale
        
        if best_met is None or best_corr < threshold:
            print(f"\n  Stopped at iteration {iteration}: best correlation {best_corr:.4f} < threshold")
            break
        
        # Calculate concentration
        concentration = best_scale * ref_concs[best_met]
        
        # Only accept if concentration is reasonable (> 0.1 mM)
        if concentration < 0.1:
            print(f"  Skipping {best_met}: concentration {concentration:.2f} mM too low")
            continue
        
        # Store result
        results[best_met] = {
            'concentration': concentration,
            'scale': best_scale,
            'correlation': best_corr,
            'iteration': iteration
        }
        
        # Subtract contribution
        contribution = best_scale * pure_spectra[best_met]
        residual -= contribution
        
        print(f"  Iter {iteration:2d}: {best_met:<15} conc={concentration:6.2f} mM  "
              f"corr={best_corr:.3f}  scale={best_scale:.3f}")
    
    # Calculate final R²
    fitted = np.zeros_like(mixture)
    for met_name, result in results.items():
        fitted += result['scale'] * pure_spectra[met_name]
    
    r2 = r2_score(mixture, fitted)
    residual_norm = np.linalg.norm(residual)
    
    print(f"\n  Final R²: {r2:.4f}")
    print(f"  Residual norm: {residual_norm:.4f}")
    
    # Results summary
    print(f"\n{'='*70}")
    print("Iterative Subtraction Results - File 10")
    print(f"{'='*70}")
    print(f"{'Metabolite':<15} {'Conc (mM)':>12} {'Iteration':>10}")
    print("-"*70)
    
    total = 0
    for met_name in sorted(results.keys()):
        conc = results[met_name]['concentration']
        iter_num = results[met_name]['iteration']
        total += conc
        print(f"{met_name:<15} {conc:>12.2f} {iter_num:>10}")
    
    print(f"{'='*70}")
    print(f"{'Total':<15} {total:>12.2f}")
    print(f"{'='*70}")
    
    # Plot
    fig, axes = plt.subplots(3, 1, figsize=(16, 12))
    
    # Original and fit
    ax = axes[0]
    ax.plot(common_ppm, mixture, 'b-', linewidth=0.8, label='Mixture', alpha=0.7)
    ax.plot(common_ppm, fitted, 'r--', linewidth=1.0, label='Iterative Fit')
    ax.set_xlim(5.5, 0.5)
    ax.set_xlabel('ppm')
    ax.set_ylabel('Normalized Intensity')
    ax.set_title(f'File 10 - Iterative Subtraction (R² = {r2:.4f})')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Residual
    ax = axes[1]
    ax.plot(common_ppm, residual, 'g-', linewidth=0.8)
    ax.axhline(y=0, color='k', linestyle='--', alpha=0.5)
    ax.set_xlim(5.5, 0.5)
    ax.set_xlabel('ppm')
    ax.set_ylabel('Residual')
    ax.set_title('Residual after subtraction')
    ax.grid(True, alpha=0.3)
    
    # Individual contributions (stacked)
    ax = axes[2]
    cumulative = np.zeros_like(mixture)
    colors = plt.cm.tab20(np.linspace(0, 1, len(results)))
    
    for i, (met_name, result) in enumerate(sorted(results.items(), key=lambda x: x[1]['concentration'], reverse=True)):
        contribution = result['scale'] * pure_spectra[met_name]
        ax.fill_between(common_ppm, cumulative, cumulative + contribution, 
                        label=f"{met_name} ({result['concentration']:.1f} mM)",
                        alpha=0.7, color=colors[i])
        cumulative += contribution
    
    ax.plot(common_ppm, mixture, 'k-', linewidth=1.0, label='Original', alpha=0.5)
    ax.set_xlim(5.5, 0.5)
    ax.set_xlabel('ppm')
    ax.set_ylabel('Cumulative Intensity')
    ax.set_title('Individual Metabolite Contributions')
    ax.legend(loc='upper right', fontsize=8, ncol=2)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('iterative_subtraction_file10.png', dpi=150, bbox_inches='tight')
    print("\nSaved: iterative_subtraction_file10.png")
    plt.close()
    
    # Compare with Lorentzian
    print(f"\n{'='*70}")
    print("Comparison with Lorentzian Results")
    print(f"{'='*70}")
    
    lorentzian = {
        'Alanine': 7.68, 'Arginine': 7.20, 'Asparagine': 1.13, 'Aspartate': 0.53,
        'Glucose': 28.98, 'Glutamate': 4.31, 'Glutamine': 5.03, 'Isoleucine': 3.25,
        'Lactate': 27.81, 'Leucine': 8.98, 'Methionine': 1.22,
        'Phenylalanine': 5.99, 'Tyrosine': 0.76, 'Valine': 0.61
    }
    
    print(f"{'Metabolite':<15} {'Lorentzian':>12} {'Iterative':>12} {'Ratio I/L':>12}")
    print("-"*70)
    
    for met in sorted(lorentzian.keys()):
        lor = lorentzian[met]
        itr = results.get(met, {}).get('concentration', 0)
        ratio = itr / lor if lor > 0 else 0
        marker = "***" if 0.5 < ratio < 2.0 else ""
        print(f"{met:<15} {lor:>12.2f} {itr:>12.2f} {ratio:>12.2f} {marker}")
    
    return results, r2

if __name__ == "__main__":
    iterative_subtraction_quantification()