physics_region_detector.py

#!/usr/bin/env python3
"""
Physics-Based Region Detection for NMR Quantification

Principles:
1. PURE-SHIFT NMR: Each chemically distinct environment → 1 singlet peak
2. DIASTEREOTOPIC: Adjacent to chiral center → 2 peaks for CH2
3. REGION GROUPING: Peaks within 0.3 ppm are grouped into single region
4. BOUNDARIES: Determined at 10% of maximum peak height
"""

import numpy as np
from scipy.signal import find_peaks
from scipy.ndimage import gaussian_filter1d


def group_expected_peaks(expected_peaks_info, merge_threshold=0.3):
    """
    Group expected peaks that are close together into regions.
    
    Parameters:
    -----------
    expected_peaks_info : list of tuples [(expected_ppm, name, protons), ...]
    merge_threshold : float
        Maximum distance (ppm) to group peaks into same region
        
    Returns:
    --------
    grouped_regions : list of dict
        Each dict contains grouped peak information
    """
    if not expected_peaks_info:
        return []
    
    # Sort by chemical shift
    sorted_peaks = sorted(expected_peaks_info, key=lambda x: x[0])
    
    grouped = []
    current_group = {
        'peaks': [sorted_peaks[0]],
        'min_ppm': sorted_peaks[0][0],
        'max_ppm': sorted_peaks[0][0],
        'total_protons': sorted_peaks[0][2]
    }
    
    for peak in sorted_peaks[1:]:
        ppm, name, protons = peak
        
        # Check if this peak is close to the current group
        if abs(ppm - current_group['max_ppm']) <= merge_threshold:
            # Add to current group
            current_group['peaks'].append(peak)
            current_group['max_ppm'] = ppm
            current_group['total_protons'] += protons
        else:
            # Start new group
            grouped.append(current_group)
            current_group = {
                'peaks': [peak],
                'min_ppm': ppm,
                'max_ppm': ppm,
                'total_protons': protons
            }
    
    # Add last group
    grouped.append(current_group)
    
    # Create region info
    regions = []
    for group in grouped:
        region = {
            'expected_centers': [p[0] for p in group['peaks']],
            'names': [p[1] for p in group['peaks']],
            'name': ' + '.join([p[1] for p in group['peaks']]),
            'protons': group['total_protons'],
            'expected_n_peaks': len(group['peaks']),
            'search_center': (group['min_ppm'] + group['max_ppm']) / 2,
            'search_width': max(0.4, group['max_ppm'] - group['min_ppm'] + 0.2)
        }
        regions.append(region)
    
    return regions


def detect_actual_region(ppm, intensity, region_info, 
                         height_threshold=0.1, 
                         margin_points=10,
                         max_width=0.5):
    """
    Detect actual region boundaries from spectrum based on physics info.
    
    Parameters:
    -----------
    ppm : array
        Chemical shift axis
    intensity : array
        Intensity values
    region_info : dict
        From group_expected_peaks()
    height_threshold : float
        Fraction of max height for boundary detection (default 0.1 = 10%)
    margin_points : int
        Additional points to add as margin beyond 10% height
    max_width : float
        Maximum region width in ppm
        
    Returns:
    --------
    detected : dict
        Detected region information
    """
    center = region_info['search_center']
    width = region_info['search_width']
    
    # Extract search window
    mask = (ppm >= center - width/2) & (ppm <= center + width/2)
    x = ppm[mask]
    y = intensity[mask]
    
    if len(x) == 0 or np.max(y) <= 0:
        return None
    
    # Smooth for peak detection
    y_smooth = gaussian_filter1d(y, sigma=2)
    
    # Find peaks
    max_y = np.max(y_smooth)
    min_height = max_y * height_threshold
    
    peaks_idx, properties = find_peaks(y_smooth, height=min_height, distance=5)
    
    if len(peaks_idx) == 0:
        # No peaks found - use expected position
        return {
            'name': region_info['name'],
            'bounds': (center - 0.1, center + 0.1),
            'peaks': [center],
            'protons': region_info['protons'],
            'expected_n_peaks': region_info['expected_n_peaks'],
            'detected_n_peaks': 0
        }
    
    # Get detected peak positions
    peak_positions = x[peaks_idx]
    peak_heights = y_smooth[peaks_idx]
    
    # Match to expected number of peaks
    # Take top N peaks where N = expected_n_peaks (or all if fewer)
    n_expected = region_info['expected_n_peaks']
    if len(peak_positions) > n_expected:
        # Sort by height and take top N
        sorted_idx = np.argsort(peak_heights)[::-1]
        selected_peaks = peak_positions[sorted_idx[:n_expected]]
    else:
        selected_peaks = peak_positions
    
    # Find boundaries at 10% height for the full region
    # Find all points above threshold
    above_threshold = y_smooth > min_height
    indices_above = np.where(above_threshold)[0]
    
    if len(indices_above) > 0:
        left_idx = max(0, indices_above[0] - margin_points)
        right_idx = min(len(x) - 1, indices_above[-1] + margin_points)
        left_bound = x[left_idx]
        right_bound = x[right_idx]
        region_start = min(left_bound, right_bound)
        region_end = max(left_bound, right_bound)
    else:
        # Fallback
        region_start = np.min(selected_peaks) - 0.05
        region_end = np.max(selected_peaks) + 0.05
    
    # Apply max width constraint
    current_width = region_end - region_start
    if current_width > max_width:
        region_center = (region_start + region_end) / 2
        region_start = region_center - max_width / 2
        region_end = region_center + max_width / 2
    
    return {
        'name': region_info['name'],
        'bounds': (region_start, region_end),
        'peaks': list(selected_peaks),
        'protons': region_info['protons'],
        'expected_n_peaks': n_expected,
        'detected_n_peaks': len(peak_positions)
    }


def get_physics_based_regions(metabolite_name, ppm, intensity, expected_peaks):
    """
    Main function: Get physics-based regions for a metabolite.
    
    Parameters:
    -----------
    metabolite_name : str
        Name of metabolite (for logging)
    ppm : array
        Chemical shift axis (TSP-corrected)
    intensity : array
        Intensity values
    expected_peaks : list
        List of (ppm, name, protons) tuples
        
    Returns:
    --------
    detected_regions : list of dict
        Each dict contains detected region information
    """
    # Group expected peaks into regions
    grouped = group_expected_peaks(expected_peaks, merge_threshold=0.3)
    
    # Detect actual regions from spectrum
    detected_regions = []
    for region_info in grouped:
        detected = detect_actual_region(ppm, intensity, region_info)
        if detected:
            detected_regions.append(detected)
    
    return detected_regions


# Physics-based expected peaks for each metabolite
# Format: (expected_ppm, proton_name, n_protons)
# Note: In pure-shift NMR, each distinct proton environment gives exactly 1 peak
# Diastereotopic protons (CH2 next to chiral center) give 2 peaks

METABOLITE_PHYSICS = {
    'Alanine': [
        (1.48, 'CH3', 3),
        (3.78, 'α-CH', 1),
    ],
    'Arginine': [
        (1.67, 'γ-CH2', 2),
        (1.90, 'β-CH2', 2),
        (3.25, 'δ-CH2', 2),
        (3.75, 'α-CH', 1),
    ],
    'Asparagine': [
        # β-CH2 is diastereotopic (adjacent to chiral center) → splits into 2 peaks
        (2.75, 'β-CH2a (diastereotopic)', 1),
        (2.85, 'β-CH2b (diastereotopic)', 1),
        (3.98, 'α-CH', 1),
    ],
    'Aspartate': [
        (2.67, 'β-CH2', 2),
        (3.90, 'α-CH', 1),
    ],
    'Glutamate': [
        (2.05, 'γ-CH2', 2),
        (2.35, 'β-CH2', 2),
        (3.75, 'α-CH', 1),
    ],
    'Glutamine': [
        (2.12, 'β-CH2', 2),
        (2.46, 'γ-CH2', 2),
        (3.78, 'α-CH', 1),
    ],
    'Isoleucine': [
        (0.95, 'δ-CH3', 3),
        (1.02, 'γ-CH3', 3),
        (1.26, 'γ-CH2', 2),
        (1.98, 'β-CH', 1),
        (3.68, 'α-CH', 1),
    ],
    'Lactate': [
        (1.32, 'CH3', 3),
        (4.12, 'CH', 1),
    ],
    'Leucine': [
        (0.97, 'δ-CH3', 3),
        (0.99, 'δ-CH3', 3),
        (1.65, 'γ-CH', 1),
        (1.72, 'β-CH2', 2),
        (3.74, 'α-CH', 1),
    ],
    'Valine': [
        (0.99, 'CH3', 3),
        (1.04, 'CH3', 3),
        (2.28, 'β-CH', 1),
        (3.62, 'α-CH', 1),
    ],
    'Phenylalanine': [
        (7.33, 'Ar-H (meta)', 2),
        (7.38, 'Ar-H (para)', 1),
        (7.43, 'Ar-H (ortho)', 2),
        (3.12, 'β-CH2', 2),
        (3.98, 'α-CH', 1),
    ],
    'Tyrosine': [
        (6.90, 'Ar-H (meta)', 2),
        (7.20, 'Ar-H (ortho)', 2),
        (3.05, 'β-CH2', 2),
        (3.95, 'α-CH', 1),
    ],
    'Methionine': [
        (2.15, 'CH3', 3),
        (2.65, 'γ-CH2', 2),
    ],
    'Glucose': [
        (5.24, 'H-1 (anomeric)', 1),
    ],
}


if __name__ == '__main__':
    import nmrglue as ng
    
    print('='*90)
    print('PHYSICS-BASED REGION DETECTION - TEST')
    print('='*90)
    print()
    
    for met_name, expected in METABOLITE_PHYSICS.items():
        filepath = f'raw_data/Reference_Raw_Date_JCAMP-DX/{met_name}-Reference/10.dx'
        
        try:
            dic, data_list = ng.jcampdx.read(filepath)
            data = data_list[0]
            
            # Get ppm scale
            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(data))
            
            # TSP correction
            mask_tsp = (ppm >= -0.5) & (ppm <= 0.5)
            tsp_ppm = ppm[mask_tsp][np.argmax(data[mask_tsp])]
            ppm_corr = ppm - tsp_ppm
            
            # Detect regions
            regions = get_physics_based_regions(met_name, ppm_corr, data, expected)
            
            print(f'{met_name}:')
            print('-'*90)
            for i, reg in enumerate(regions):
                print(f'  Region {i+1}: {reg["name"]}')
                print(f'    Bounds: {reg["bounds"][0]:.3f} - {reg["bounds"][1]:.3f} ppm '
                      f'(width: {reg["bounds"][1] - reg["bounds"][0]:.3f})')
                print(f'    Peaks: {[f"{p:.3f}" for p in reg["peaks"]]}')
                print(f'    Protons: {reg["protons"]}')
            print()
            
        except Exception as e:
            print(f'{met_name}: Error - {e}')
            print()