Merge pull request #7 from catalystneuro/stub-test-duration-sec

add stub_test_duration_sec parameter to ks 2.5
catalystneuro · Feb 6, 2024 · a2810c7 · a2810c7
2 parents 7f819fb + 93e04c6
commit a2810c7
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/common/models.py b/common/models.py
@@ -72,6 +72,7 @@ class PipelineContext(BaseModel):
     lazy_read_input: bool = Field(default=True, description='Lazy read input file')
     write_recording: bool = Field(default=False, description='Write recording')
     stub_test: bool = Field(default=False, description='Stub test')
+    stub_test_duration_sec: float = Field(default=300, description='Stub test duration in seconds')
     job_kwargs: JobKwargs = Field(default=JobKwargs(), description='Job kwargs')
     recording_context: RecordingContext = Field(description='Recording context')
     run_preprocessing: bool = Field(default=True, description='Run preprocessing')

diff --git a/common/processor_pipeline.py b/common/processor_pipeline.py
@@ -44,7 +44,8 @@ def run_pipeline(context: PipelineContext):
 
     if context.stub_test:
         logger.info('Running in stub test mode')
-        n_frames = int(min(300_000, recording.get_num_frames()))
+        stub_test_num_frames = context.stub_test_duration_sec * recording.get_sampling_frequency()
+        n_frames = int(min(stub_test_num_frames, recording.get_num_frames()))
         recording = recording.frame_slice(start_frame=0, end_frame=n_frames)
 
     logger.info(recording)

diff --git a/si_kilosort25/spec.json b/si_kilosort25/spec.json
@@ -15,12 +15,14 @@
                     "description": "Input NWB file"
                 }
             ],
+            "inputFolders": [],
             "outputs": [
                 {
                     "name": "output",
                     "description": "Output NWB file"
                 }
             ],
+            "outputFolders": [],
             "parameters": [
                 {
                     "name": "lazy_read_input",
@@ -40,6 +42,12 @@
                     "type": "bool",
                     "default": false
                 },
+                {
+                    "name": "stub_test_duration_sec",
+                    "description": "Stub test duration in seconds",
+                    "type": "float",
+                    "default": 300
+                },
                 {
                     "name": "job_kwargs.n_jobs",
                     "description": "Number of jobs, must be a positive number between 0 and 1, or -1 for all processors.",