plot_tsp_corrected_spectra.py

#!/usr/bin/env python3
"""
TSP-Corrected NMR Spectra Plotting Script
=========================================

This script:
1. Finds the TSP peak in each dx file
2. Corrects the ppm axis so TSP is at exactly 0 ppm
3. Plots the aliphatic region (0.9-3.5 ppm) with 0.05 ppm minor ticks

Author: Assistant
Date: 2026-03-23
"""

import nmrglue as ng
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, AutoMinorLocator
from scipy.ndimage import gaussian_filter1d
from pathlib import Path
import re


def get_even_numbered_files(data_dir):
    """Get all even-numbered .dx files from the specified directory."""
    dx_files = sorted(data_dir.glob("*.dx"), 
                      key=lambda x: int(re.search(r'(\d+)', x.name).group()))
    even_files = [f for f in dx_files 
                  if int(re.search(r'(\d+)', f.name).group()) % 2 == 0]
    return even_files


def read_nmr_data(file_path):
    """
    Read NMR data from JCAMP-DX file and calculate magnitude spectrum.
    Returns ppm axis and magnitude spectrum (NOT corrected for TSP).
    """
    # Read JCAMP-DX file
    dic, data_list = ng.jcampdx.read(str(file_path))
    
    # Get real and imaginary parts
    data_real = data_list[0]
    data_imag = data_list[1] if len(data_list) > 1 else np.zeros_like(data_real)
    
    # Calculate magnitude spectrum
    data_magnitude = np.sqrt(data_real**2 + data_imag**2)
    
    # Calculate ppm axis
    sfo1 = float(dic['$SFO1'][0])  # Spectrometer frequency (MHz)
    o1_hz = float(dic['$O1'][0])   # Offset frequency (Hz)
    sw_hz = float(dic['$SWH'][0])  # Spectral width (Hz)
    si = data_real.size            # Number of points
    
    o1_ppm = o1_hz / sfo1          # Offset in ppm
    sw_ppm = sw_hz / sfo1          # Spectral width in ppm
    
    # Create ppm axis (left to right: high ppm -> low ppm)
    ppm = np.linspace(o1_ppm + sw_ppm/2, o1_ppm - sw_ppm/2, si)
    
    return ppm, data_magnitude, sfo1


def find_tsp_peak(ppm, data, tsp_region=(-0.5, 0.5)):
    """
    Find the TSP peak position in the specified region.
    Returns the ppm position of the TSP peak maximum.
    """
    mask = (ppm >= tsp_region[0]) & (ppm <= tsp_region[1])
    
    if not np.any(mask):
        return None
    
    ppm_region = ppm[mask]
    data_region = data[mask]
    
    # Find the maximum peak
    peak_idx = np.argmax(data_region)
    peak_ppm = ppm_region[peak_idx]
    peak_intensity = data_region[peak_idx]
    
    # Calculate noise for S/N check
    noise = np.std(data_region[data_region < np.percentile(data_region, 50)])
    sn_ratio = peak_intensity / noise if noise > 0 else 0
    
    # Only return if S/N is reasonable (>10)
    if sn_ratio < 10:
        return None
    
    return peak_ppm


def process_file_with_tsp_correction(file_path, sigma=2):
    """
    Process a single file: read data, find TSP, correct ppm axis.
    
    Returns:
        dict with 'ppm_corrected', 'smoothed', 'tsp_original', 'tsp_corrected'
    """
    # Read raw data
    ppm, magnitude, sfo1 = read_nmr_data(file_path)
    
    # Apply Gaussian smoothing
    smoothed = gaussian_filter1d(magnitude, sigma=sigma)
    
    # Find TSP peak in original ppm axis
    tsp_ppm = find_tsp_peak(ppm, smoothed)
    
    if tsp_ppm is None:
        # If TSP not found, use default correction of -0.0784 ppm
        # (from our statistical analysis)
        tsp_ppm = -0.0784
        print(f"  ⚠️ TSP not found in {file_path.name}, using default correction")
    
    # Calculate correction: we want TSP to be at 0 ppm
    # If TSP is at -0.08 ppm, we need to ADD 0.08 ppm to all values
    correction = -tsp_ppm  # e.g., if tsp_ppm = -0.08, correction = +0.08
    
    # Apply correction to ppm axis
    ppm_corrected = ppm + correction
    
    return {
        'ppm': ppm_corrected,
        'magnitude': magnitude,
        'smoothed': smoothed,
        'tsp_original': tsp_ppm,
        'tsp_corrected': 0.0,  # After correction, TSP is at 0
        'correction_applied': correction,
        'sfo1': sfo1
    }


def plot_tsp_corrected_spectra(data_dir, output_file=None, region=(0.9, 3.5), 
                                sigma=2, show_tsp_region=True):
    """
    Plot TSP-corrected spectra for all files in a folder.
    
    Parameters:
    -----------
    data_dir : Path
        Directory containing .dx files
    output_file : str or Path, optional
        Output filename. If None, uses folder name.
    region : tuple
        PPM region to plot (default: 0.9-3.5 ppm, aliphatic)
    sigma : float
        Gaussian smoothing parameter
    show_tsp_region : bool
        Whether to include a small subplot showing the TSP region
    """
    even_files = get_even_numbered_files(data_dir)
    folder_name = data_dir.name.replace('-Reference', '')
    
    if output_file is None:
        output_file = f"{folder_name.lower()}_tsp_corrected.png"
    
    # Color scheme - use tab10 for distinct colors
    colors = plt.cm.tab10(np.linspace(0, 1, len(even_files)))
    
    # Create figure with main plot and optional TSP inset
    if show_tsp_region:
        fig = plt.figure(figsize=(16, 10))
        gs = fig.add_gridspec(2, 2, height_ratios=[3, 1], hspace=0.3, wspace=0.3)
        ax_main = fig.add_subplot(gs[0, :])  # Main plot spans both columns
        ax_tsp = fig.add_subplot(gs[1, 0])   # TSP region
        ax_stats = fig.add_subplot(gs[1, 1])  # Statistics
    else:
        fig, ax_main = plt.subplots(figsize=(14, 8))
    
    tsp_positions = []
    corrections = []
    
    # Plot each spectrum
    for i, file_path in enumerate(even_files):
        try:
            # Process with TSP correction
            data = process_file_with_tsp_correction(file_path, sigma=sigma)
            
            # Store TSP info
            tsp_positions.append(data['tsp_original'])
            corrections.append(data['correction_applied'])
            
            # Extract region for plotting
            mask = (data['ppm'] >= region[0]) & (data['ppm'] <= region[1])
            ppm_plot = data['ppm'][mask]
            intensity_plot = data['smoothed'][mask]
            
            # Normalize intensity for better visualization (optional)
            # intensity_plot = intensity_plot / np.max(intensity_plot)
            
            # Plot
            ax_main.plot(ppm_plot, intensity_plot, 
                        color=colors[i], linewidth=1.5, 
                        label=file_path.stem, alpha=0.8)
            
            # Plot TSP region if requested
            if show_tsp_region:
                tsp_mask = (data['ppm'] >= -0.5) & (data['ppm'] <= 0.5)
                if np.any(tsp_mask):
                    ax_tsp.plot(data['ppm'][tsp_mask], data['smoothed'][tsp_mask],
                               color=colors[i], linewidth=1.5, alpha=0.7)
                
        except Exception as e:
            print(f"Error processing {file_path.name}: {e}")
    
    # Configure main plot
    ax_main.set_xlabel('Chemical Shift (ppm)', fontsize=14, fontweight='bold')
    ax_main.set_ylabel('Intensity (a.u.)', fontsize=14, fontweight='bold')
    ax_main.set_title(f'¹H NMR Spectra - {folder_name} (TSP-Corrected)\n'
                     f'Aliphatic Region ({region[0]}-{region[1]} ppm), '
                     f'Gaussian smoothing (σ={sigma})',
                     fontsize=14, fontweight='bold')
    
    # Reverse x-axis (standard NMR convention: high ppm to low ppm)
    ax_main.set_xlim(region[1], region[0])
    
    # Major and minor ticks - 0.05 ppm minor ticks
    ax_main.xaxis.set_major_locator(MultipleLocator(0.5))   # Major: every 0.5 ppm
    ax_main.xaxis.set_minor_locator(MultipleLocator(0.05))  # Minor: every 0.05 ppm
    ax_main.tick_params(which='major', length=6, width=1.5, direction='out')
    ax_main.tick_params(which='minor', length=3, width=1, direction='out')
    
    # Grid
    ax_main.grid(True, which='major', alpha=0.3, linestyle='--')
    ax_main.grid(True, which='minor', alpha=0.15, linestyle=':')
    
    # Legend
    ax_main.legend(loc='upper right', fontsize=9, title='File Number', 
                  ncol=2 if len(even_files) > 8 else 1)
    
    # Configure TSP region subplot
    if show_tsp_region:
        ax_tsp.set_xlabel('Chemical Shift (ppm)', fontsize=11)
        ax_tsp.set_ylabel('Intensity', fontsize=11)
        ax_tsp.set_title('TSP Reference Region (-0.5 to 0.5 ppm)', fontsize=12, fontweight='bold')
        ax_tsp.set_xlim(0.5, -0.5)  # Reversed
        ax_tsp.axvline(x=0, color='red', linestyle='--', linewidth=2, alpha=0.7, label='TSP at 0 ppm')
        ax_tsp.xaxis.set_major_locator(MultipleLocator(0.2))
        ax_tsp.xaxis.set_minor_locator(MultipleLocator(0.05))
        ax_tsp.tick_params(which='major', length=5, direction='out')
        ax_tsp.tick_params(which='minor', length=3, direction='out')
        ax_tsp.grid(True, alpha=0.3)
        ax_tsp.legend(loc='upper right')
        
        # Statistics subplot
        ax_stats.axis('off')
        stats_text = (
            f"TSP Correction Statistics\n"
            f"{'='*30}\n"
            f"Files processed: {len(tsp_positions)}\n"
            f"Mean TSP position: {np.mean(tsp_positions):.4f} ppm\n"
            f"Std deviation: {np.std(tsp_positions):.4f} ppm\n"
            f"Mean correction: {np.mean(corrections):+.4f} ppm\n"
            f"Correction range: {np.min(corrections):+.4f} to {np.max(corrections):+.4f} ppm"
        )
        ax_stats.text(0.1, 0.5, stats_text, fontsize=11, family='monospace',
                     verticalalignment='center')
    
    plt.tight_layout()
    plt.savefig(output_file, dpi=200, bbox_inches='tight')
    print(f"✅ TSP-corrected plot saved to: {output_file}")
    plt.close()
    
    return {
        'folder': folder_name,
        'tsp_positions': tsp_positions,
        'corrections': corrections,
        'mean_correction': np.mean(corrections)
    }


def process_all_folders(base_dir, output_dir, region=(0.9, 3.5), sigma=2):
    """
    Process all metabolite folders and generate TSP-corrected plots.
    """
    folders = sorted([d for d in base_dir.iterdir() if d.is_dir()])
    
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)
    
    all_stats = []
    
    print("=" * 70)
    print("TSP-CORRECTED SPECTRA GENERATION")
    print("=" * 70)
    print(f"\nBase directory: {base_dir}")
    print(f"Output directory: {output_dir}")
    print(f"Plot region: {region[0]}-{region[1]} ppm")
    print(f"Gaussian smoothing: σ={sigma}")
    print(f"Minor ticks: 0.05 ppm")
    print(f"\nFound {len(folders)} folders to process")
    
    for folder in folders:
        print(f"\n{'='*70}")
        print(f"Processing: {folder.name}")
        print(f"{'='*70}")
        
        output_file = output_dir / f"{folder.name.lower().replace('-reference', '')}_tsp_corrected.png"
        
        try:
            stats = plot_tsp_corrected_spectra(
                folder, 
                output_file=str(output_file),
                region=region,
                sigma=sigma,
                show_tsp_region=True
            )
            all_stats.append(stats)
        except Exception as e:
            print(f"Error processing {folder.name}: {e}")
    
    # Print summary
    print(f"\n{'='*70}")
    print("SUMMARY - TSP CORRECTIONS APPLIED")
    print(f"{'='*70}")
    print(f"\n{'Metabolite':<25} {'Mean Correction (ppm)':<20} {'Files'}")
    print("-" * 70)
    for s in all_stats:
        print(f"{s['folder']:<25} {s['mean_correction']:>+18.4f}  {len(s['tsp_positions']):>3}")
    
    print(f"\n{'='*70}")
    print("All TSP-corrected plots generated!")
    print(f"Output directory: {output_dir}")
    print(f"{'='*70}")
    
    return all_stats


def main():
    """
    Main function to generate TSP-corrected spectra for all metabolites.
    """
    # Configuration
    BASE_DIR = Path("/home/tjiang/elise/pure_shift_nmr/new/raw_data/"
                    "Reference_Raw_Date_JCAMP-DX")
    OUTPUT_DIR = Path("/home/tjiang/elise/pure_shift_nmr/new/tsp_corrected_spectra")
    
    # Plotting parameters
    REGION = (0.9, 3.5)  # Aliphatic region
    SIGMA = 2            # Gaussian smoothing
    
    # Process all folders
    process_all_folders(BASE_DIR, OUTPUT_DIR, region=REGION, sigma=SIGMA)


if __name__ == "__main__":
    main()