Skip to content

Commit 4194e0c

Browse files
alejoe91alejoe91
committed
Add script to compile and build and test
1 parent 6bb034f commit 4194e0c

2 files changed

Lines changed: 237 additions & 0 deletions

File tree

‎kilosort-compiled/compile_and_build_all_ks.sh‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
#!/bin/bash
2+
set -e
3+
4+
if [ $# == 0 ]; then
5+
echo "Usage: $0 param1 param2 param3 param4 param4"
6+
echo "* param1: spikeinterface path"
7+
echo "* param1: kilosort1 path"
8+
echo "* param1: kilosort2 path"
9+
echo "* param1: kilosort2.5 path"
10+
echo "* param1: kilosort3 path"
11+
exit
12+
fi
13+
14+
if [ $# -ne 5 ]; then
15+
echo "spikeinterface and kilosort paths (1, 2, 2.5, 3) must be given"
16+
exit 1
17+
fi
18+
19+
20+
SPIKEINTERFACE_PATH=${1%/}
21+
KILOSORT_PATH=${2%/}
22+
KILOSORT2_PATH=${3%/}
23+
KILOSORT25_PATH=${4%/}
24+
KILOSORT3_PATH=${5%/}
25+
26+
27+
cd kilosort-compiled
28+
bash compile.sh $KILOSORT_PATH $SPIKEINTERFACE_PATH
29+
bash build.sh
30+
cd ..
31+
32+
cd kilosort2-compiled
33+
bash compile.sh $KILOSORT2_PATH $SPIKEINTERFACE_PATH
34+
bash build.sh
35+
cd ..
36+
37+
cd kilosort2_5-compiled
38+
bash compile.sh $KILOSORT25_PATH $SPIKEINTERFACE_PATH
39+
bash build.sh
40+
cd ..
41+
42+
cd kilosort3-compiled
43+
bash compile.sh $KILOSORT3_PATH $SPIKEINTERFACE_PATH
44+
bash build.sh
45+
cd ..

‎tests/test_kilosorts.py‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
import os
2+
import shutil
3+
from pathlib import Path
4+
5+
import pytest
6+
7+
import spikeinterface as si
8+
import spikeinterface.extractors as se
9+
import spikeinterface.preprocessing as spre
10+
import spikeinterface.sorters as ss
11+
12+
13+
si_path = Path(si.__file__)
14+
15+
if "site-package" not in si_path.name:
16+
print("Using SPIKEINTERFACE_DEV_PATH")
17+
SI_DEV_PATH = str(si_path.parent.parent.parent)
18+
else:
19+
SI_DEV_PATH = None
20+
21+
os.environ["SINGULARITY_DISABLE_CACHE"] = "true"
22+
23+
# test docker or singularity
24+
DOCKER_SINGULARITY = "docker" # "singularity" | "docker"
25+
26+
27+
def generate_run_kwargs():
28+
test_recording, _ = se.toy_example(duration=30, seed=0, num_channels=64, num_segments=1)
29+
cache_folder = Path("cached_recording")
30+
if cache_folder.is_dir():
31+
shutil.rmtree(cache_folder)
32+
test_recording = test_recording.save(folder="cached_recording")
33+
test_recording.set_channel_gains(1)
34+
test_recording.set_channel_offsets(0)
35+
run_kwargs = dict(recording=test_recording, verbose=True)
36+
if DOCKER_SINGULARITY == "singularity":
37+
run_kwargs["singularity_image"] = True
38+
elif DOCKER_SINGULARITY == "docker":
39+
run_kwargs["docker_image"] = True
40+
else:
41+
raise Exception("DOCKER_SINGULARITY can be 'docker' or 'singularity'")
42+
return run_kwargs
43+
44+
45+
@pytest.fixture(autouse=True)
46+
def work_dir(request, tmp_path):
47+
"""
48+
This fixture, along with "run_kwargs" creates one folder per
49+
test function using built-in tmp_path pytest fixture
50+
51+
The tmp_path will be the working directory for the test function
52+
53+
At the end of the each test function, a clean up will be done
54+
"""
55+
os.chdir(tmp_path)
56+
yield
57+
os.chdir(request.config.invocation_dir)
58+
shutil.rmtree(str(tmp_path))
59+
60+
61+
@pytest.fixture
62+
def run_kwargs(work_dir):
63+
return generate_run_kwargs()
64+
65+
66+
def test_kilosort(run_kwargs):
67+
import os
68+
69+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
70+
sorting = ss.run_sorter("kilosort", output_folder="kilosort", **run_kwargs)
71+
print(sorting)
72+
73+
74+
def test_kilosort2(run_kwargs):
75+
import os
76+
77+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
78+
sorting = ss.run_sorter("kilosort2", output_folder="kilosort2", **run_kwargs)
79+
print(sorting)
80+
81+
82+
def test_kilosort2_release_si(run_kwargs):
83+
# docker image
84+
if DOCKER_SINGULARITY == "docker":
85+
run_kwargs["docker_image"] = "spikeinterface/kilosort2_5-compiled-base:0.2.0"
86+
else:
87+
run_kwargs["singularity_image"] = "spikeinterface/kilosort2_5-compiled-base:0.2.0"
88+
import os
89+
90+
del os.environ["SPIKEINTERFACE_DEV_PATH"]
91+
sorting = ss.run_sorter("kilosort2", output_folder="kilosort2", **run_kwargs)
92+
print(sorting)
93+
94+
95+
def test_kilosort2_5(run_kwargs):
96+
import os
97+
98+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
99+
sorting = ss.run_sorter("kilosort2_5", output_folder="kilosort2_5", **run_kwargs)
100+
print(sorting)
101+
102+
103+
def test_kilosort3(run_kwargs):
104+
import os
105+
106+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
107+
sorting = ss.run_sorter("kilosort3", output_folder="kilosort3", **run_kwargs)
108+
print(sorting)
109+
110+
111+
def test_kilosort2_skip_preproc(run_kwargs):
112+
import os
113+
114+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
115+
print("Skipping preprocessing")
116+
gain = 200
117+
run_kwargs["recording"] = spre.zscore(run_kwargs["recording"])
118+
run_kwargs["recording"] = spre.scale(run_kwargs["recording"], gain=gain, dtype="int16")
119+
120+
sorting = ss.run_sorter(
121+
"kilosort2", output_folder="kilosort2", skip_kilosort_preprocessing=True, scaleproc=gain, **run_kwargs
122+
)
123+
print(sorting)
124+
125+
126+
def test_kilosort2_5_skip_preproc(run_kwargs):
127+
import os
128+
129+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
130+
print("Skipping preprocessing")
131+
gain = 200
132+
run_kwargs["recording"] = spre.zscore(run_kwargs["recording"])
133+
run_kwargs["recording"] = spre.scale(run_kwargs["recording"], gain=gain, dtype="int16")
134+
sorting = ss.run_sorter(
135+
"kilosort2_5", output_folder="kilosort2_5", skip_kilosort_preprocessing=True, scaleproc=gain, **run_kwargs
136+
)
137+
print(sorting)
138+
139+
140+
def test_kilosort2_5_skip_preproc_old_image(run_kwargs):
141+
import os
142+
143+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
144+
gain = 200
145+
run_kwargs["recording"] = spre.zscore(run_kwargs["recording"])
146+
run_kwargs["recording"] = spre.scale(run_kwargs["recording"], gain=gain, dtype="int16")
147+
148+
# docker image
149+
if DOCKER_SINGULARITY == "docker":
150+
run_kwargs["docker_image"] = "spikeinterface/kilosort2_5-compiled-base:0.1.0"
151+
else:
152+
run_kwargs["singularity_image"] = "spikeinterface/kilosort2_5-compiled-base:0.1.0"
153+
print(run_kwargs)
154+
sorting = ss.run_sorter(
155+
"kilosort2_5", output_folder="kilosort2_5", skip_kilosort_preprocessing=True, scaleproc=gain, **run_kwargs
156+
)
157+
print(sorting)
158+
159+
160+
def test_kilosort3_skip_preproc(run_kwargs):
161+
import os
162+
163+
os.environ["SPIKEINTERFACE_DEV_PATH"] = SI_DEV_PATH
164+
print("Skipping preprocessing")
165+
run_kwargs["recording"] = spre.zscore(run_kwargs["recording"])
166+
run_kwargs["recording"] = spre.scale(run_kwargs["recording"], gain=200, dtype="int16")
167+
sorting = ss.run_sorter(
168+
"kilosort3", output_folder="kilosort3", skip_kilosort_preprocessing=True, scaleproc=200, **run_kwargs
169+
)
170+
print(sorting)
171+
172+
173+
if __name__ == "__main__":
174+
kwargs = generate_run_kwargs()
175+
print("\n\nKilosort")
176+
test_kilosort(kwargs)
177+
178+
print("\n\nKilosort2")
179+
test_kilosort2(kwargs)
180+
test_kilosort2_skip_preproc(kwargs)
181+
print("\nKilosort2 released si")
182+
test_kilosort2_release_si(kwargs)
183+
184+
print("\n\nKilosort2.5")
185+
test_kilosort2_5(kwargs)
186+
test_kilosort2_5_skip_preproc(kwargs)
187+
print("\nKilosort2.5 old image")
188+
test_kilosort2_5_skip_preproc_old_image(kwargs)
189+
190+
print("\n\nKilosort3")
191+
test_kilosort3(kwargs)
192+
test_kilosort3_skip_preproc(kwargs)

0 commit comments

Comments
 (0)