1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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()
|
