path: root/examples/example_diffusion.py

#!/usr/bin/env python3
"""
Example demonstrating diffusion processing of transient events.

This script shows how to use the new event_processor module to analyze
diffusion characteristics of detected events, similar to the approach
in heehun_diffusion_processing.py.
"""

import numpy as np
import matplotlib.pyplot as plt

# Import transivent functions
from transivent import configure_logging, detect

# Import building blocks from submodules
from transivent.analysis import (
    detect_events,
    calculate_initial_background,
    estimate_noise,
    analyze_thresholds,
)
from transivent.event_processor import (
    process_events_for_diffusion,
    extract_event_waveforms,
    calculate_msd_parallel,
    calculate_acf,
    fit_diffusion_linear,
    plot_diffusion_comparison,
)
configure_logging("INFO")


def create_synthetic_dataset(
    n_events: int = 50,
    event_duration: float = 100e-6,  # 100 microseconds
    sampling_interval: float = 1e-6,  # 1 microsecond
    total_duration: float = 0.1,  # 100 milliseconds
    diffusion_coeff: float = 1e-12,  # m²/s
    noise_level: float = 0.01,  # Lower noise for better detection
    signal_amplitude: float = 2.0,  # Higher amplitude for better detection
):
    """
    Create synthetic data with diffusion-like events.
    
    Parameters
    ----------
    n_events : int
        Number of events to generate
    event_duration : float
        Duration of each event in seconds
    sampling_interval : float
        Sampling interval in seconds
    total_duration : float
        Total duration of the signal in seconds
    diffusion_coeff : float
        Diffusion coefficient for event dynamics
    noise_level : float
        Noise level in the signal
    signal_amplitude : float
        Amplitude of the events
        
    Returns
    -------
    t : np.ndarray
        Time array
    x : np.ndarray
        Signal array
    """
    n_points = int(total_duration / sampling_interval)
    t = np.linspace(0, total_duration, n_points)
    x = np.random.normal(0, noise_level, n_points)
    
    # Generate random event times
    event_starts = np.random.uniform(0, total_duration - event_duration, n_events)
    
    for start_time in event_starts:
        start_idx = int(start_time / sampling_interval)
        end_idx = int((start_time + event_duration) / sampling_interval)
        
        if end_idx >= n_points:
            end_idx = n_points - 1
        
        # Create a diffusion-like signal for this event
        event_length = end_idx - start_idx
        event_t = np.linspace(0, event_duration, event_length)
        
        # Simulate Brownian motion
        np.random.seed(int(start_time * 1e6))  # Seed for reproducibility
        steps = np.random.normal(0, np.sqrt(2 * diffusion_coeff * sampling_interval), event_length)
        event_signal = signal_amplitude * (1 + np.cumsum(steps))
        
        # Add to main signal
        x[start_idx:end_idx] += event_signal
    
    return t, x


def analyze_single_dataset(name: str, t: np.ndarray, x: np.ndarray, sampling_interval: float):
    """
    Analyze a single dataset for diffusion characteristics.
    
    Parameters
    ----------
    name : str
        Name of the dataset
    t : np.ndarray
        Time array
    x : np.ndarray
        Signal array
    sampling_interval : float
        Sampling interval in seconds
        
    Returns
    -------
    dict
        Results containing events and diffusion analysis
    """
    print(f"\n=== Analyzing {name} ===")
    
    # Step 1: Detect events using transivent
    # First calculate background
    smooth_n = int(10e-6 / sampling_interval)  # 10 microsecond smoothing window
    if smooth_n % 2 == 0:
        smooth_n += 1
    
    bg_initial = calculate_initial_background(t, x, smooth_n)
    global_noise = estimate_noise(x, bg_initial)
    
    # Detect events
    events, _ = detect_events(
        t, x, bg_initial,
        snr_threshold=3.0,
        min_event_len=int(10e-6 / sampling_interval),
        min_event_amp=5.0 * global_noise,
        widen_frac=0.5,
        global_noise=global_noise,
        signal_polarity=1,  # Positive events
    )
    
    print(f"Detected {len(events)} events")
    
    # Step 2: Process events for diffusion analysis
    results = process_events_for_diffusion(
        name=name,
        sampling_interval=sampling_interval,
        data_path="",
        t=t,
        x=x,
        events=events,
        max_lag=500,
        n_jobs=-1,
    )
    
    print(f"Processed {results['event_count']} events for diffusion")
    if results['event_count'] > 0:
        print(f"Mean diffusion coefficient: {results['statistics']['mean_diffusion']:.3e} ± "
              f"{results['statistics']['std_diffusion']:.3e} m²/s")
    
    return results


def main():
    """Main function demonstrating diffusion analysis."""
    print("Transivent Diffusion Processing Example")
    print("=" * 40)
    
    # Create synthetic datasets with different diffusion characteristics
    datasets = {
        "Ferritin (slow diffusion)": {
            "diffusion_coeff": 5e-13,  # Slower diffusion
            "signal_amplitude": 1.0,
            "n_events": 40,
        },
        "Catalase (medium diffusion)": {
            "diffusion_coeff": 1e-12,  # Medium diffusion
            "signal_amplitude": 0.8,
            "n_events": 35,
        },
        "BSA (fast diffusion)": {
            "diffusion_coeff": 2e-12,  # Faster diffusion
            "signal_amplitude": 0.6,
            "n_events": 45,
        },
    }
    
    sampling_interval = 1e-6  # 1 microsecond
    
    all_results = {}
    
    # Analyze each dataset
    for name, params in datasets.items():
        # Generate synthetic data
        t, x = create_synthetic_dataset(
            n_events=params["n_events"],
            diffusion_coeff=params["diffusion_coeff"],
            signal_amplitude=params["signal_amplitude"],
            sampling_interval=sampling_interval,
        )
        
        # Analyze for diffusion
        results = analyze_single_dataset(name, t, x, sampling_interval)
        all_results[name] = results
    
    # Create comparison plot
    if any(len(results["diffusion_coeffs"]) > 0 for results in all_results.values()):
        print("\n=== Creating comparison plot ===")
        fig = plot_diffusion_comparison(
            all_results,
            show_ellipses=True,
            colors=["#1f77b4", "gold", "green"],
            markers=["o", "^", "s"],
        )
        
        # Save the plot
        import os
        os.makedirs("analysis_output", exist_ok=True)
        fig.savefig("analysis_output/diffusion_comparison.png", dpi=150, bbox_inches="tight")
        print("Saved plot to: analysis_output/diffusion_comparison.png")
        
        # Show plot
        plt.show()
    else:
        print("\nNo events detected in any dataset. Check detection parameters.")


if __name__ == "__main__":
    main()