diff options
Diffstat (limited to 'tests/test_event_processor.py')
| -rw-r--r-- | tests/test_event_processor.py | 174 |
1 files changed, 174 insertions, 0 deletions
diff --git a/tests/test_event_processor.py b/tests/test_event_processor.py new file mode 100644 index 0000000..cd23334 --- /dev/null +++ b/tests/test_event_processor.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +"""Test script for the event_processor module.""" + +import numpy as np +import matplotlib.pyplot as plt + +# Test the event_processor module +from transivent import ( + extract_event_waveforms, + calculate_msd_parallel, + calculate_acf, + fit_diffusion_linear, + process_events_for_diffusion, +) + +def test_extract_event_waveforms(): + """Test event waveform extraction.""" + print("Testing extract_event_waveforms...") + + # Create test data + t = np.linspace(0, 1, 1000) + x = np.random.normal(0, 0.1, 1000) + + # Add some events + x[200:300] += 1.0 # Event 1 + x[500:600] += 0.5 # Event 2 + + # Define events + events = np.array([[0.2, 0.3], [0.5, 0.6]]) + + # Extract waveforms + waveforms = extract_event_waveforms(t, x, events) + + assert len(waveforms) == 2, f"Expected 2 waveforms, got {len(waveforms)}" + assert len(waveforms[0]) == 100, f"Expected 100 points in first event, got {len(waveforms[0])}" + assert len(waveforms[1]) == 100, f"Expected 100 points in second event, got {len(waveforms[1])}" + + # Check that events have been extracted correctly + assert np.mean(waveforms[0]) > 0.5, "First event should have positive mean" + assert np.mean(waveforms[1]) > 0.2, "Second event should have positive mean" + + print("✓ extract_event_waveforms test passed") + + +def test_calculate_msd(): + """Test MSD calculation.""" + print("Testing calculate_msd_parallel...") + + # Create a simple Brownian motion trajectory + n_points = 1000 + dt = 1e-6 + D_true = 1e-12 # True diffusion coefficient + + # Generate random walk + np.random.seed(42) + steps = np.random.normal(0, np.sqrt(2 * D_true * dt), n_points) + x = np.cumsum(steps) + + # Calculate MSD + taus, msds, counts = calculate_msd_parallel(x, dt=dt, max_lag=100, n_jobs=1) + + assert len(taus) == 100, f"Expected 100 lag times, got {len(taus)}" + assert len(msds) == 100, f"Expected 100 MSD values, got {len(msds)}" + assert len(counts) == 100, f"Expected 100 counts, got {len(counts)}" + + # Check that MSD increases with time (for diffusive motion) + assert msds[50] > msds[10], "MSD should increase with lag time" + + print("✓ calculate_msd_parallel test passed") + + +def test_calculate_acf(): + """Test ACF calculation.""" + print("Testing calculate_acf...") + + # Create a simple signal + n_points = 1000 + dt = 1e-6 + x = np.random.normal(0, 1, n_points) + + # Add some correlation + for i in range(1, n_points): + x[i] += 0.5 * x[i-1] + + # Calculate ACF + lags, acf = calculate_acf(x, dt=dt, max_lag=100) + + assert len(lags) == 101, f"Expected 101 lag values (0 to 100), got {len(lags)}" + assert len(acf) == 101, f"Expected 101 ACF values, got {len(acf)}" + + # Check that ACF at lag 0 is variance + assert np.abs(acf[0] - np.var(x)) < 0.1, f"ACF(0) should be variance, got {acf[0]}" + + # Check that ACF decreases with lag (for correlated signal) + assert acf[50] < acf[10], "ACF should decrease with lag time" + + print("✓ calculate_acf test passed") + + +def test_fit_diffusion(): + """Test diffusion coefficient fitting.""" + print("Testing fit_diffusion_linear...") + + # Create synthetic MSD data + taus = np.linspace(0, 1e-5, 100) + D_true = 1e-12 + msds = 2 * D_true * taus + 0.1 * np.random.normal(0, 1e-24, 100) # Add noise + + # Fit diffusion + D_fit = fit_diffusion_linear(taus, msds, time_limit=3e-5) + + assert not np.isnan(D_fit), "Diffusion coefficient should not be NaN" + assert D_fit > 0, "Diffusion coefficient should be positive" + + # Check that fitted value is close to true value (within 50% due to noise) + assert np.abs(D_fit - D_true) / D_true < 0.5, f"Fitted D ({D_fit}) too far from true D ({D_true})" + + print("✓ fit_diffusion_linear test passed") + + +def test_process_events_for_diffusion(): + """Test the high-level wrapper function.""" + print("Testing process_events_for_diffusion...") + + # Create test data + t = np.linspace(0, 1, 1000) + x = np.random.normal(0, 0.1, 1000) + + # Add some events + x[200:300] += 1.0 # Event 1 + x[500:600] += 0.5 # Event 2 + + # Define events + events = np.array([[0.2, 0.3], [0.5, 0.6]]) + + # Process events + results = process_events_for_diffusion( + name="test", + sampling_interval=1e-3, # 1 ms + data_path="", + t=t, + x=x, + events=events, + max_lag=50, + n_jobs=1, + ) + + assert "diffusion_coeffs" in results, "Results should contain diffusion_coeffs" + assert "acf_values" in results, "Results should contain acf_values" + assert "event_count" in results, "Results should contain event_count" + assert "statistics" in results, "Results should contain statistics" + + assert len(results["diffusion_coeffs"]) == 2, f"Expected 2 diffusion coeffs, got {len(results['diffusion_coeffs'])}" + assert len(results["acf_values"]) == 2, f"Expected 2 ACF values, got {len(results['acf_values'])}" + assert results["event_count"] == 2, f"Expected event_count=2, got {results['event_count']}" + + print("✓ process_events_for_diffusion test passed") + + +def main(): + """Run all tests.""" + print("Running event_processor module tests...\n") + + test_extract_event_waveforms() + test_calculate_msd() + test_calculate_acf() + test_fit_diffusion() + test_process_events_for_diffusion() + + print("\n✅ All tests passed!") + + +if __name__ == "__main__": + main()
\ No newline at end of file |