From 93e04c69ace2ccf8490f0f764f38a596195d0fca Mon Sep 17 00:00:00 2001
From: Jeremy Magland <jmagland@flatironinstitute.org>
Date: Tue, 6 Feb 2024 11:19:55 -0500
Subject: [PATCH] add stub_test_duration_sec

---
 common/models.py             | 1 +
 common/processor_pipeline.py | 3 ++-
 si_kilosort25/spec.json      | 8 ++++++++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/common/models.py b/common/models.py
index 00054c5..efe21b9 100644
--- 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
index 5c4289b..c7f477c 100644
--- 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
index a97e70c..a2f5128 100644
--- 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.",