test_region_wise_lasso_v2.py

#!/usr/bin/env python3
"""
Region-wise LASSO v2 - Fixed normalization and constraints
Key fixes:
1. Use TSP-normalized spectra (not max-normalized) for proper concentration scaling
2. Ensure metabolites only appear in their primary regions
3. Use stricter LASSO regularization
4. Sum concentrations from different regions, don't average
"""

import sys
sys.path.insert(0, '.')
import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import Lasso, LassoCV, ElasticNetCV
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
)

# Table S6 values for M2 from SI paper
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.42, 'Valine': 0.33, 'Tyrosine': 0.33  # Estimated Met/Val/Tyr
}

class RegionWiseLASSOv2:
    def __init__(self, reference_base_dir="raw_data/Reference_Raw_Date_JCAMP-DX"):
        self.reference_dir = reference_base_dir
        self.metabolite_info = get_metabolite_info_v3()
        self.pure_spectra = {}
        
        # Define regions with specific metabolites
        # Only include metabolites that have PRIMARY signals in each region
        self.regions = {
            'aliphatic': {
                'bounds': (0.5, 1.8),
                'metabolites': ['Alanine', 'Valine', 'Arginine', 'Lactate', 
                               'Isoleucine', 'Leucine', 'Methionine', 'Glutamate']
            },
            'methine': {
                'bounds': (1.8, 2.6),
                'metabolites': ['Glutamate', 'Aspartate', 'Glutamine', 
                               'Asparagine', 'Isoleucine', 'Leucine', 'Methionine']
            },
            'ch_oh': {
                'bounds': (3.0, 4.5),
                'metabolites': ['Glucose', 'Lactate', 'Alanine', 'Asparagine',
                               'Aspartate', 'Glutamine', 'Glutamate', 'Methionine']
            },
            'aromatic': {
                'bounds': (6.5, 7.8),
                'metabolites': ['Phenylalanine', 'Tyrosine']
            },
        }
        
        self._load_references()
    
    def _load_references(self):
        """Load TSP-normalized reference spectra"""
        print("Loading reference spectra (TSP-normalized)...")
        
        for met_name, info in self.metabolite_info.items():
            ref_file = os.path.join(self.reference_dir, 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))
                
                # CRITICAL: Use TSP-normalized spectrum, NOT max-normalized
                # This preserves the relative concentration information
                spec_norm = spec / tsp_area
                
                self.pure_spectra[met_name] = {
                    'ppm': ppm_corr,
                    'spectrum': spec_norm,
                    'concentration': info['files'][20]
                }
                
            except Exception as e:
                print(f"  Error loading {met_name}: {e}")
        
        print(f"  Loaded {len(self.pure_spectra)} references")
    
    def extract_region(self, ppm, spec, region):
        """Extract spectral region"""
        mask = (ppm >= region[0]) & (ppm <= region[1])
        return ppm[mask], spec[mask]
    
    def build_region_dictionary(self, region_info, target_ppm):
        """Build dictionary for a specific region with only relevant metabolites"""
        bounds = region_info['bounds']
        allowed_mets = region_info['metabolites']
        
        dictionary = []
        labels = []
        ref_concs = []
        
        for met_name in allowed_mets:
            if met_name not in self.pure_spectra:
                continue
                
            data = self.pure_spectra[met_name]
            ppm, spec = data['ppm'], data['spectrum']
            
            # Extract region
            ppm_reg, spec_reg = self.extract_region(ppm, spec, bounds)
            
            if len(ppm_reg) == 0:
                continue
            
            # Interpolate to target ppm grid
            spec_interp = np.interp(target_ppm, ppm_reg[::-1], spec_reg[::-1])
            
            # Only include if there's significant signal in this region
            if np.max(spec_interp) > 0.001:  # Lower threshold
                dictionary.append(spec_interp)
                labels.append(met_name)
                ref_concs.append(data['concentration'])
        
        if len(dictionary) == 0:
            return None, None, None
        
        X = np.column_stack(dictionary)
        return X, labels, np.array(ref_concs)
    
    def quantify_region(self, region_name, region_info, mixture_ppm, mixture_spec):
        """Quantify metabolites in a single region using LASSO"""
        
        print(f"\n{'='*60}")
        print(f"Region: {region_name} ({region_info['bounds'][0]:.1f}-{region_info['bounds'][1]:.1f} ppm)")
        print(f"Allowed: {region_info['metabolites']}")
        print(f"{'='*60}")
        
        bounds = region_info['bounds']
        
        # Extract mixture region
        ppm_mix, spec_mix = self.extract_region(mixture_ppm, mixture_spec, bounds)
        
        if len(ppm_mix) == 0:
            print("  No data in region")
            return {}
        
        # Create common ppm grid for this region
        target_ppm = np.linspace(bounds[0], bounds[1], 400)
        spec_mix_interp = np.interp(target_ppm, ppm_mix[::-1], spec_mix[::-1])
        
        # Build dictionary
        X, labels, ref_concs = self.build_region_dictionary(region_info, target_ppm)
        
        if X is None or len(labels) == 0:
            print("  No reference spectra in region")
            return {}
        
        print(f"  Dictionary: {X.shape[1]} metabolites")
        print(f"  {labels}")
        
        # LASSO with cross-validation - use stricter alpha range
        print(f"  Running LASSO...")
        
        # Normalize columns of X for stable fitting
        X_norm = X.copy()
        scales = []
        for i in range(X.shape[1]):
            col_max = np.max(np.abs(X[:, i]))
            if col_max > 0:
                X_norm[:, i] = X[:, i] / col_max
                scales.append(col_max)
            else:
                scales.append(1.0)
        scales = np.array(scales)
        
        model = LassoCV(cv=3, fit_intercept=False, positive=True,
                       alphas=np.logspace(-5, -2, 40), max_iter=10000)
        
        try:
            model.fit(X_norm, spec_mix_interp)
            
            # Predict and calculate R²
            y_pred = model.predict(X_norm)
            r2 = r2_score(spec_mix_interp, y_pred)
            
            print(f"  Optimal alpha: {model.alpha_:.6f}")
            print(f"  R²: {r2:.4f}")
            
            # Extract results
            results = {}
            for i, (label, coef) in enumerate(zip(labels, model.coef_)):
                if coef > 0.001:  # Threshold
                    # De-normalize coefficient
                    actual_coef = coef / scales[i]
                    
                    # Calculate concentration: coef * ref_conc
                    # The coefficient represents the scaling factor relative to reference
                    conc = actual_coef * ref_concs[i]
                    
                    results[label] = {
                        'concentration': conc,
                        'coefficient': coef,
                        'actual_coef': actual_coef,
                        'ref_conc': ref_concs[i]
                    }
                    print(f"    {label}: {conc:.2f} mM (coef={coef:.4f})")
            
            return results
            
        except Exception as e:
            print(f"  Error: {e}")
            import traceback
            traceback.print_exc()
            return {}
    
    def quantify_mixture(self, mixture_file):
        """Quantify mixture using region-wise LASSO"""
        
        print(f"{'='*70}")
        print(f"Region-wise LASSO v2 - FIXED")
        print(f"Mixture: {os.path.basename(mixture_file)}")
        print(f"{'='*70}")
        
        # 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}")
        
        # Quantify each region
        all_results = {}
        
        for region_name, region_info in self.regions.items():
            region_results = self.quantify_region(
                region_name, region_info, ppm_corr, spec_norm
            )
            all_results[region_name] = region_results
        
        # Combine results - use SMART aggregation
        # For metabolites appearing in multiple regions, use the one with highest confidence
        combined = {}
        region_source = {}
        
        for region_name, results in all_results.items():
            for met_name, data in results.items():
                conc = data['concentration']
                coef = data['coefficient']
                
                # Keep the result with highest coefficient (most confident fit)
                if met_name not in combined or coef > combined[met_name]['coef']:
                    combined[met_name] = {'conc': conc, 'coef': coef}
                    region_source[met_name] = region_name
        
        # Final summary
        print(f"\n{'='*70}")
        print("FINAL RESULTS - Region-wise LASSO v2")
        print(f"{'='*70}")
        print(f"{'Metabolite':<15} {'Concentration (mM)':>20} {'Region':>15}")
        print("-"*70)
        
        total = 0
        for met_name in sorted(combined.keys()):
            conc = combined[met_name]['conc']
            total += conc
            print(f"{met_name:<15} {conc:>20.2f} {region_source[met_name]:>15}")
        
        print(f"{'='*70}")
        print(f"{'Total':<15} {total:>20.2f}")
        print(f"{'='*70}")
        
        return combined, all_results


def main():
    """Test region-wise LASSO v2 on File 10 (which is M2)"""
    
    quantifier = RegionWiseLASSOv2()
    
    mixture_file = "raw_data/Model_Mixtures-M1-M6_JCAMP-DX/10.dx"
    results, region_details = quantifier.quantify_mixture(mixture_file)
    
    # Compare with Table S6 (M2 values)
    print(f"\n{'='*70}")
    print("Comparison with Table S6 (M2 reference values from paper)")
    print(f"{'='*70}")
    
    print(f"{'Metabolite':<15} {'Table S6':>12} {'LASSO v2':>12} {'Ratio':>10}")
    print("-"*70)
    
    for met in sorted(TABLE_S6_M2.keys()):
        expected = TABLE_S6_M2[met]
        observed = results.get(met, {}).get('conc', 0) if isinstance(results.get(met), dict) else results.get(met, 0)
        if isinstance(observed, dict):
            observed = observed.get('conc', 0)
        
        ratio = observed / expected if expected > 0 else 0
        
        # Mark if within 20%
        marker = "OK" if 0.8 <= ratio <= 1.2 else ""
        print(f"{met:<15} {expected:>12.2f} {observed:>12.2f} {ratio:>10.2f} {marker}")


if __name__ == "__main__":
    main()