nmr_integration_script.py

#!/usr/bin/env python3
"""
NMR Peak Integration and Plotting Script
========================================

This script reads JCAMP-DX NMR files, calculates magnitude spectra,
applies Gaussian smoothing, integrates peaks in specified regions,
and generates comparison plots.

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

import nmrglue as ng
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend for saving plots
import matplotlib.pyplot as plt
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.
    
    Parameters:
    -----------
    data_dir : Path
        Directory containing the .dx files
    
    Returns:
    --------
    list : Sorted list of even-numbered file paths
    """
    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.
    
    Parameters:
    -----------
    file_path : Path or str
        Path to the .dx file
    
    Returns:
    --------
    dict : Dictionary containing:
        - 'ppm': ppm axis
        - 'magnitude': magnitude spectrum
        - 'smoothed': smoothed magnitude spectrum
        - 'params': NMR parameters (SFO1, O1, SWH, etc.)
    """
    # 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 (always positive)
    data_magnitude = np.sqrt(data_real**2 + data_imag**2)
    
    # Apply Gaussian smoothing (sigma=3)
    data_smoothed = gaussian_filter1d(data_magnitude, sigma=3)
    
    # 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': ppm,
        'magnitude': data_magnitude,
        'smoothed': data_smoothed,
        'params': {
            'sfo1': sfo1,
            'o1_ppm': o1_ppm,
            'sw_ppm': sw_ppm,
            'si': si
        }
    }


def integrate_peak(ppm, intensity, region):
    """
    Integrate peak area within specified ppm region.
    
    Parameters:
    -----------
    ppm : array
        PPM axis
    intensity : array
        Intensity values
    region : tuple (min_ppm, max_ppm)
        Integration region boundaries
    
    Returns:
    --------
    float : Integrated area (always positive)
    """
    mask = (ppm >= region[0]) & (ppm <= region[1])
    ppm_region = ppm[mask]
    intensity_region = intensity[mask]
    
    # Use absolute value to ensure positive area
    # (ppm axis goes high->low, so trapz gives negative)
    area = abs(np.trapz(intensity_region, ppm_region))
    return area


def process_all_files(data_dir, ch3_region=(1.1, 1.8), ch_region=(3.5, 3.9)):
    """
    Process all even-numbered files and calculate integrations.
    
    Parameters:
    -----------
    data_dir : Path
        Directory containing .dx files
    ch3_region : tuple
        PPM region for CH3 peak (default: 1.1-1.8 ppm)
    ch_region : tuple
        PPM region for CH peak (default: 3.5-3.9 ppm)
    
    Returns:
    --------
    list : List of dictionaries with integration results
    """
    even_files = get_even_numbered_files(data_dir)
    results = []
    
    print(f"Found {len(even_files)} even-numbered files")
    print("=" * 70)
    
    for file_path in even_files:
        try:
            # Read data
            data = read_nmr_data(file_path)
            
            # Integrate original data
            ch3_orig = integrate_peak(data['ppm'], data['magnitude'], ch3_region)
            ch_orig = integrate_peak(data['ppm'], data['magnitude'], ch_region)
            
            # Integrate smoothed data
            ch3_smooth = integrate_peak(data['ppm'], data['smoothed'], ch3_region)
            ch_smooth = integrate_peak(data['ppm'], data['smoothed'], ch_region)
            
            # Calculate ratios
            ratio_orig = ch3_orig / ch_orig if ch_orig > 0 else 0
            ratio_smooth = ch3_smooth / ch_smooth if ch_smooth > 0 else 0
            
            results.append({
                'file': file_path.stem,
                'ch3_original': ch3_orig,
                'ch3_smoothed': ch3_smooth,
                'ch_original': ch_orig,
                'ch_smoothed': ch_smooth,
                'ratio_orig': ratio_orig,
                'ratio_smooth': ratio_smooth
            })
            
            print(f"\n{file_path.stem}:")
            print(f"  CH3 ({ch3_region[0]}-{ch3_region[1]} ppm):  "
                  f"Original = {ch3_orig:.2e}, Smoothed = {ch3_smooth:.2e}")
            print(f"  CH  ({ch_region[0]}-{ch_region[1]} ppm):  "
                  f"Original = {ch_orig:.2e}, Smoothed = {ch_smooth:.2e}")
            print(f"  CH3/CH Ratio:       Original = {ratio_orig:.2f}, "
                  f"Smoothed = {ratio_smooth:.2f}")
            
        except Exception as e:
            print(f"Error processing {file_path.name}: {e}")
    
    return results


def plot_integration_results(results, output_file='integration_results.png'):
    """
    Create bar plots comparing original and smoothed integration results.
    
    Parameters:
    -----------
    results : list
        List of integration result dictionaries
    output_file : str
        Output filename for the plot
    """
    files = [r['file'] for r in results]
    x = np.arange(len(files))
    width = 0.35
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # Plot 1: CH3 region comparison
    ax1 = axes[0, 0]
    ch3_orig = [r['ch3_original'] for r in results]
    ch3_smooth = [r['ch3_smoothed'] for r in results]
    ax1.bar(x - width/2, ch3_orig, width, label='Original', 
            color='skyblue', edgecolor='black')
    ax1.bar(x + width/2, ch3_smooth, width, label='Smoothed (σ=3)', 
            color='lightcoral', edgecolor='black')
    ax1.set_xlabel('File Number')
    ax1.set_ylabel('Integration Area')
    ax1.set_title('CH₃ Region - Peak Area')
    ax1.set_xticks(x)
    ax1.set_xticklabels(files)
    ax1.legend()
    ax1.grid(True, alpha=0.3, axis='y')
    
    # Plot 2: CH region comparison
    ax2 = axes[0, 1]
    ch_orig = [r['ch_original'] for r in results]
    ch_smooth = [r['ch_smoothed'] for r in results]
    ax2.bar(x - width/2, ch_orig, width, label='Original', 
            color='skyblue', edgecolor='black')
    ax2.bar(x + width/2, ch_smooth, width, label='Smoothed (σ=3)', 
            color='lightcoral', edgecolor='black')
    ax2.set_xlabel('File Number')
    ax2.set_ylabel('Integration Area')
    ax2.set_title('CH Region - Peak Area')
    ax2.set_xticks(x)
    ax2.set_xticklabels(files)
    ax2.legend()
    ax2.grid(True, alpha=0.3, axis='y')
    
    # Plot 3: CH3/CH ratio comparison
    ax3 = axes[1, 0]
    ratios_orig = [r['ratio_orig'] for r in results]
    ratios_smooth = [r['ratio_smooth'] for r in results]
    ax3.bar(x - width/2, ratios_orig, width, label='Original', 
            color='lightgreen', edgecolor='black')
    ax3.bar(x + width/2, ratios_smooth, width, label='Smoothed (σ=3)', 
            color='orange', edgecolor='black')
    ax3.set_xlabel('File Number')
    ax3.set_ylabel('CH₃/CH Area Ratio')
    ax3.set_title('Peak Area Ratio (CH₃/CH)')
    ax3.set_xticks(x)
    ax3.set_xticklabels(files)
    ax3.axhline(y=3, color='r', linestyle='--', alpha=0.5, 
                label='Theoretical (3:1)')
    ax3.legend()
    ax3.grid(True, alpha=0.3, axis='y')
    
    # Plot 4: Deviation from theoretical ratio
    ax4 = axes[1, 1]
    theoretical = 3.0
    deviations_orig = [abs(r['ratio_orig'] - theoretical) / theoretical * 100 
                       for r in results]
    deviations_smooth = [abs(r['ratio_smooth'] - theoretical) / theoretical * 100 
                         for r in results]
    ax4.bar(x - width/2, deviations_orig, width, label='Original', 
            color='purple', edgecolor='black', alpha=0.7)
    ax4.bar(x + width/2, deviations_smooth, width, label='Smoothed (σ=3)', 
            color='brown', edgecolor='black', alpha=0.7)
    ax4.set_xlabel('File Number')
    ax4.set_ylabel('Deviation from Theoretical (%)')
    ax4.set_title('Deviation from Theoretical Ratio (3:1)')
    ax4.set_xticks(x)
    ax4.set_xticklabels(files)
    ax4.legend()
    ax4.grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig(output_file, dpi=200, bbox_inches='tight')
    print(f"\n✅ Integration plot saved to: {output_file}")
    plt.close()


def plot_spectra_overlay(data_dir, output_file='spectra_overlay.png', 
                         region=(0.5, 5.0)):
    """
    Plot all spectra overlaid in one figure.
    
    Parameters:
    -----------
    data_dir : Path
        Directory containing .dx files
    output_file : str
        Output filename for the plot
    region : tuple
        PPM region to plot (default: 0.5-5.0 ppm)
    """
    even_files = get_even_numbered_files(data_dir)
    colors = ['blue', 'red', 'green', 'purple', 'orange', 'brown', 'magenta', 'cyan']
    
    fig, ax = plt.subplots(figsize=(14, 8))
    
    for i, file_path in enumerate(even_files):
        try:
            data = read_nmr_data(file_path)
            mask = (data['ppm'] >= region[0]) & (data['ppm'] <= region[1])
            color = colors[i % len(colors)]
            
            ax.plot(data['ppm'][mask], data['smoothed'][mask], 
                   color=color, linewidth=1.5, label=file_path.stem, alpha=0.8)
        except Exception as e:
            print(f"Error plotting {file_path.name}: {e}")
    
    ax.set_xlabel('Chemical Shift (ppm)', fontsize=14, fontweight='bold')
    ax.set_ylabel('Intensity (a.u.)', fontsize=14, fontweight='bold')
    ax.set_title(f'¹H NMR Spectra - Alanine (Even-numbered files)\n'
                 f'Smoothed (σ=3), {region[0]}-{region[1]} ppm', 
                 fontsize=16, fontweight='bold')
    ax.set_xlim(region[1], region[0])  # Reverse x-axis
    ax.legend(loc='upper right', fontsize=10, title='File Number')
    ax.grid(True, alpha=0.3, linestyle='--')
    
    plt.tight_layout()
    plt.savefig(output_file, dpi=200, bbox_inches='tight')
    print(f"✅ Spectra overlay saved to: {output_file}")
    plt.close()


def print_summary_table(results):
    """
    Print a formatted summary table of integration results.
    
    Parameters:
    -----------
    results : list
        List of integration result dictionaries
    """
    print("\n" + "=" * 95)
    print("SUMMARY TABLE")
    print("=" * 95)
    print(f"{'File':<8} {'CH3 Original':<15} {'CH3 Smoothed':<15} "
          f"{'CH Original':<15} {'CH Smoothed':<15} {'Ratio Orig':<12} "
          f"{'Ratio Smooth':<12}")
    print("-" * 95)
    
    for r in results:
        print(f"{r['file']:<8} {r['ch3_original']:<15.4e} "
              f"{r['ch3_smoothed']:<15.4e} {r['ch_original']:<15.4e} "
              f"{r['ch_smoothed']:<15.4e} {r['ratio_orig']:<12.2f} "
              f"{r['ratio_smooth']:<12.2f}")


def main():
    """
    Main function to run the integration and plotting workflow.
    """
    # Configuration
    DATA_DIR = Path("/home/tjiang/elise/pure_shift_nmr/new/raw_data/"
                    "Reference_Raw_Date_JCAMP-DX/Alanine-Reference")
    
    # Integration regions (ppm)
    CH3_REGION = (1.1, 1.8)   # CH3 doublet region
    CH_REGION = (3.5, 3.9)    # CH quartet region
    
    print("=" * 70)
    print("NMR PEAK INTEGRATION AND PLOTTING")
    print("=" * 70)
    print(f"\nData directory: {DATA_DIR}")
    print(f"CH3 region: {CH3_REGION[0]}-{CH3_REGION[1]} ppm")
    print(f"CH region: {CH_REGION[0]}-{CH_REGION[1]} ppm")
    print(f"Smoothing: Gaussian σ=3")
    
    # Process all files
    results = process_all_files(DATA_DIR, CH3_REGION, CH_REGION)
    
    # Print summary table
    print_summary_table(results)
    
    # Generate plots
    print("\n" + "=" * 70)
    print("GENERATING PLOTS")
    print("=" * 70)
    
    plot_integration_results(results, 
                             output_file='integration_results.png')
    
    plot_spectra_overlay(DATA_DIR, 
                         output_file='spectra_overlay.png',
                         region=(0.5, 5.0))
    
    print("\n" + "=" * 70)
    print("ANALYSIS COMPLETE")
    print("=" * 70)
    print("\nNotes:")
    print("  - Integration uses trapezoidal rule (np.trapz)")
    print("  - Absolute value ensures positive areas")
    print("  - Theoretical CH3/CH ratio = 3.0 (3 protons vs 1 proton)")
    print("  - Deviations indicate experimental or processing artifacts")


if __name__ == "__main__":
    main()