Skip to content

Commit ad8436e

Browse files
authored
Merge pull request #88 from ttngu207/no-curation
overall code cleanup/improvement for more robust and optimal kilosort run
2 parents d9c3887 + fd331bd commit ad8436e

File tree

5 files changed

+98
-11
lines changed

5 files changed

+98
-11
lines changed

element_array_ephys/ephys_acute.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,9 @@ def make(self, key):
295295
raise FileNotFoundError(
296296
'No Open Ephys data found for probe insertion: {}'.format(key))
297297

298+
if not probe_data.ap_meta:
299+
raise IOError('No analog signals found - check "structure.oebin" file or "continuous" directory')
300+
298301
if probe_data.probe_model in supported_probe_types:
299302
probe_type = probe_data.probe_model
300303
electrode_query = probe.ProbeType.Electrode & {'probe_type': probe_type}

element_array_ephys/ephys_chronic.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,9 @@ def make(self, key):
242242
raise FileNotFoundError(
243243
'No Open Ephys data found for probe insertion: {}'.format(key))
244244

245+
if not probe_data.ap_meta:
246+
raise IOError('No analog signals found - check "structure.oebin" file or "continuous" directory')
247+
245248
if probe_data.probe_model in supported_probe_types:
246249
probe_type = probe_data.probe_model
247250
electrode_query = probe.ProbeType.Electrode & {'probe_type': probe_type}

element_array_ephys/ephys_no_curation.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,9 @@ def make(self, key):
293293
raise FileNotFoundError(
294294
'No Open Ephys data found for probe insertion: {}'.format(key))
295295

296+
if not probe_data.ap_meta:
297+
raise IOError('No analog signals found - check "structure.oebin" file or "continuous" directory')
298+
296299
if probe_data.probe_model in supported_probe_types:
297300
probe_type = probe_data.probe_model
298301
electrode_query = probe.ProbeType.Electrode & {'probe_type': probe_type}

element_array_ephys/readers/kilosort_triggering.py

Lines changed: 88 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import os
88
import scipy.io
99
import numpy as np
10-
from datetime import datetime
10+
from datetime import datetime, timedelta
1111

1212
from element_interface.utils import dict_to_uuid
1313

@@ -191,14 +191,17 @@ def run_modules(self):
191191
if module_status['completion_time'] is not None:
192192
continue
193193

194-
module_output_json = module_input_json.replace('-input.json',
195-
'-' + module + '-output.json')
194+
module_output_json = self._get_module_output_json_filename(module)
196195
command = (sys.executable
197196
+ " -W ignore -m ecephys_spike_sorting.modules." + module
198197
+ " --input_json " + module_input_json
199198
+ " --output_json " + module_output_json)
200199

201200
start_time = datetime.utcnow()
201+
self._update_module_status(
202+
{module: {'start_time': start_time,
203+
'completion_time': None,
204+
'duration': None}})
202205
with open(module_logfile, "a") as f:
203206
subprocess.check_call(command.split(' '), stdout=f)
204207
completion_time = datetime.utcnow()
@@ -207,6 +210,8 @@ def run_modules(self):
207210
'completion_time': completion_time,
208211
'duration': (completion_time - start_time).total_seconds()}})
209212

213+
self._update_total_duration()
214+
210215
def _get_raw_data_filepaths(self):
211216
session_str, gate_str, _, probe_str = self.parse_input_filename()
212217

@@ -248,10 +253,44 @@ def _get_module_status(self, module):
248253
if self._modules_input_hash_fp.exists():
249254
with open(self._modules_input_hash_fp) as f:
250255
modules_status = json.load(f)
256+
if modules_status[module]['completion_time'] is None:
257+
# additional logic to read from the "-output.json" file for this module as well
258+
# handle cases where the module has finished successfully,
259+
# but the "_modules_input_hash_fp" is not updated (for whatever reason),
260+
# resulting in this module not registered as completed in the "_modules_input_hash_fp"
261+
module_output_json_fp = pathlib.Path(self._get_module_output_json_filename(module))
262+
if module_output_json_fp.exists():
263+
with open(module_output_json_fp) as f:
264+
module_run_output = json.load(f)
265+
modules_status[module]['duration'] = module_run_output['execution_time']
266+
modules_status[module]['completion_time'] = (
267+
datetime.strptime(modules_status[module]['start_time'], '%Y-%m-%d %H:%M:%S.%f')
268+
+ timedelta(seconds=module_run_output['execution_time']))
251269
return modules_status[module]
252270

253271
return {'start_time': None, 'completion_time': None, 'duration': None}
254272

273+
def _get_module_output_json_filename(self, module):
274+
module_input_json = self._module_input_json.as_posix()
275+
module_output_json = module_input_json.replace(
276+
'-input.json',
277+
'-' + module + '-' + str(self._modules_input_hash) + '-output.json')
278+
return module_output_json
279+
280+
def _update_total_duration(self):
281+
with open(self._modules_input_hash_fp) as f:
282+
modules_status = json.load(f)
283+
cumulative_execution_duration = sum(
284+
v['duration'] or 0 for k, v in modules_status.items()
285+
if k not in ('cumulative_execution_duration', 'total_duration'))
286+
total_duration = (
287+
datetime.strptime(modules_status[self._modules[-1]]['completion_time'], '%Y-%m-%d %H:%M:%S.%f')
288+
- datetime.strptime(modules_status[self._modules[0]]['start_time'], '%Y-%m-%d %H:%M:%S.%f')
289+
).total_seconds()
290+
self._update_module_status(
291+
{'cumulative_execution_duration': cumulative_execution_duration,
292+
'total_duration': total_duration})
293+
255294

256295
class OpenEphysKilosortPipeline:
257296
"""
@@ -353,22 +392,27 @@ def run_modules(self):
353392
if module_status['completion_time'] is not None:
354393
continue
355394

356-
module_output_json = module_input_json.replace('-input.json',
357-
'-' + module + '-output.json')
358-
command = (sys.executable
359-
+ " -W ignore -m ecephys_spike_sorting.modules." + module
360-
+ " --input_json " + module_input_json
361-
+ " --output_json " + module_output_json)
395+
module_output_json = self._get_module_output_json_filename(module)
396+
command = [sys.executable,
397+
'-W', 'ignore', '-m', 'ecephys_spike_sorting.modules.' + module,
398+
'--input_json', module_input_json,
399+
'--output_json', module_output_json]
362400

363401
start_time = datetime.utcnow()
402+
self._update_module_status(
403+
{module: {'start_time': start_time,
404+
'completion_time': None,
405+
'duration': None}})
364406
with open(module_logfile, "a") as f:
365-
subprocess.check_call(command.split(' '), stdout=f)
407+
subprocess.check_call(command, stdout=f)
366408
completion_time = datetime.utcnow()
367409
self._update_module_status(
368410
{module: {'start_time': start_time,
369411
'completion_time': completion_time,
370412
'duration': (completion_time - start_time).total_seconds()}})
371413

414+
self._update_total_duration()
415+
372416
def _update_module_status(self, updated_module_status={}):
373417
if self._modules_input_hash is None:
374418
raise RuntimeError('"generate_modules_input_json()" not yet performed!')
@@ -393,10 +437,44 @@ def _get_module_status(self, module):
393437
if self._modules_input_hash_fp.exists():
394438
with open(self._modules_input_hash_fp) as f:
395439
modules_status = json.load(f)
440+
if modules_status[module]['completion_time'] is None:
441+
# additional logic to read from the "-output.json" file for this module as well
442+
# handle cases where the module has finished successfully,
443+
# but the "_modules_input_hash_fp" is not updated (for whatever reason),
444+
# resulting in this module not registered as completed in the "_modules_input_hash_fp"
445+
module_output_json_fp = pathlib.Path(self._get_module_output_json_filename(module))
446+
if module_output_json_fp.exists():
447+
with open(module_output_json_fp) as f:
448+
module_run_output = json.load(f)
449+
modules_status[module]['duration'] = module_run_output['execution_time']
450+
modules_status[module]['completion_time'] = (
451+
datetime.strptime(modules_status[module]['start_time'], '%Y-%m-%d %H:%M:%S.%f')
452+
+ timedelta(seconds=module_run_output['execution_time']))
396453
return modules_status[module]
397454

398455
return {'start_time': None, 'completion_time': None, 'duration': None}
399456

457+
def _get_module_output_json_filename(self, module):
458+
module_input_json = self._module_input_json.as_posix()
459+
module_output_json = module_input_json.replace(
460+
'-input.json',
461+
'-' + module + '-' + str(self._modules_input_hash) + '-output.json')
462+
return module_output_json
463+
464+
def _update_total_duration(self):
465+
with open(self._modules_input_hash_fp) as f:
466+
modules_status = json.load(f)
467+
cumulative_execution_duration = sum(
468+
v['duration'] or 0 for k, v in modules_status.items()
469+
if k not in ('cumulative_execution_duration', 'total_duration'))
470+
total_duration = (
471+
datetime.strptime(modules_status[self._modules[-1]]['completion_time'], '%Y-%m-%d %H:%M:%S.%f')
472+
- datetime.strptime(modules_status[self._modules[0]]['start_time'], '%Y-%m-%d %H:%M:%S.%f')
473+
).total_seconds()
474+
self._update_module_status(
475+
{'cumulative_execution_duration': cumulative_execution_duration,
476+
'total_duration': total_duration})
477+
400478

401479
def run_pykilosort(continuous_file, kilosort_output_directory, params,
402480
channel_ind, x_coords, y_coords, shank_ind, connected, sample_rate):

element_array_ephys/readers/openephys.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ def load_probe_data(self):
145145
else:
146146
continue # not continuous data for the current probe
147147
else:
148-
raise ValueError(f'Unable to infer type (AP or LFP) for the continuous data from:\n\t{continuous_info}')
148+
raise ValueError(f'Unable to infer type (AP or LFP) for the continuous data from:\n\t{continuous_info["folder_name"]}')
149149

150150
if continuous_type == 'ap':
151151
probe.recording_info['recording_count'] += 1

0 commit comments

Comments
 (0)