Skip to content

Commit 8ab4c58

Browse files
JaerongAJaerongA
committed
refactor: ♻️ improve if else block in EphysRecording
1 parent 673faed commit 8ab4c58

File tree

1 file changed

+118
-119
lines changed

1 file changed

+118
-119
lines changed

‎element_array_ephys/ephys_no_curation.py

Lines changed: 118 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -335,145 +335,144 @@ def make(self, key):
335335
f"Processing ephys files from acquisition software of type {acq_software} is not yet implemented."
336336
)
337337

338-
else:
339-
supported_probe_types = probe.ProbeType.fetch("probe_type")
340-
341-
if acq_software == "SpikeGLX":
342-
for meta_filepath in ephys_meta_filepaths:
343-
spikeglx_meta = spikeglx.SpikeGLXMeta(meta_filepath)
344-
if str(spikeglx_meta.probe_SN) == inserted_probe_serial_number:
345-
break
346-
else:
347-
raise FileNotFoundError(
348-
"No SpikeGLX data found for probe insertion: {}".format(key)
349-
)
338+
supported_probe_types = probe.ProbeType.fetch("probe_type")
350339

351-
if spikeglx_meta.probe_model in supported_probe_types:
352-
probe_type = spikeglx_meta.probe_model
353-
electrode_query = probe.ProbeType.Electrode & {
354-
"probe_type": probe_type
355-
}
340+
if acq_software == "SpikeGLX":
341+
for meta_filepath in ephys_meta_filepaths:
342+
spikeglx_meta = spikeglx.SpikeGLXMeta(meta_filepath)
343+
if str(spikeglx_meta.probe_SN) == inserted_probe_serial_number:
344+
break
345+
else:
346+
raise FileNotFoundError(
347+
"No SpikeGLX data found for probe insertion: {}".format(key)
348+
)
356349

357-
probe_electrodes = {
358-
(shank, shank_col, shank_row): key
359-
for key, shank, shank_col, shank_row in zip(
360-
*electrode_query.fetch(
361-
"KEY", "shank", "shank_col", "shank_row"
362-
)
363-
)
364-
}
350+
if spikeglx_meta.probe_model in supported_probe_types:
351+
probe_type = spikeglx_meta.probe_model
352+
electrode_query = probe.ProbeType.Electrode & {
353+
"probe_type": probe_type
354+
}
365355

366-
electrode_group_members = [
367-
probe_electrodes[(shank, shank_col, shank_row)]
368-
for shank, shank_col, shank_row, _ in spikeglx_meta.shankmap[
369-
"data"
370-
]
371-
]
372-
else:
373-
raise NotImplementedError(
374-
"Processing for neuropixels probe model"
375-
" {} not yet implemented".format(spikeglx_meta.probe_model)
356+
probe_electrodes = {
357+
(shank, shank_col, shank_row): key
358+
for key, shank, shank_col, shank_row in zip(
359+
*electrode_query.fetch(
360+
"KEY", "shank", "shank_col", "shank_row"
361+
)
376362
)
363+
}
377364

378-
self.insert1(
379-
{
380-
**key,
381-
**generate_electrode_config(
382-
probe_type, electrode_group_members
383-
),
384-
"acq_software": acq_software,
385-
"sampling_rate": spikeglx_meta.meta["imSampRate"],
386-
"recording_datetime": spikeglx_meta.recording_time,
387-
"recording_duration": (
388-
spikeglx_meta.recording_duration
389-
or spikeglx.retrieve_recording_duration(meta_filepath)
390-
),
391-
}
365+
electrode_group_members = [
366+
probe_electrodes[(shank, shank_col, shank_row)]
367+
for shank, shank_col, shank_row, _ in spikeglx_meta.shankmap[
368+
"data"
369+
]
370+
]
371+
else:
372+
raise NotImplementedError(
373+
"Processing for neuropixels probe model"
374+
" {} not yet implemented".format(spikeglx_meta.probe_model)
392375
)
393376

394-
root_dir = find_root_directory(get_ephys_root_data_dir(), meta_filepath)
395-
self.EphysFile.insert1(
396-
{**key, "file_path": meta_filepath.relative_to(root_dir).as_posix()}
377+
self.insert1(
378+
{
379+
**key,
380+
**generate_electrode_config(
381+
probe_type, electrode_group_members
382+
),
383+
"acq_software": acq_software,
384+
"sampling_rate": spikeglx_meta.meta["imSampRate"],
385+
"recording_datetime": spikeglx_meta.recording_time,
386+
"recording_duration": (
387+
spikeglx_meta.recording_duration
388+
or spikeglx.retrieve_recording_duration(meta_filepath)
389+
),
390+
}
391+
)
392+
393+
root_dir = find_root_directory(get_ephys_root_data_dir(), meta_filepath)
394+
self.EphysFile.insert1(
395+
{**key, "file_path": meta_filepath.relative_to(root_dir).as_posix()}
396+
)
397+
elif acq_software == "Open Ephys":
398+
dataset = openephys.OpenEphys(session_dir)
399+
for serial_number, probe_data in dataset.probes.items():
400+
if str(serial_number) == inserted_probe_serial_number:
401+
break
402+
else:
403+
raise FileNotFoundError(
404+
"No Open Ephys data found for probe insertion: {}".format(key)
397405
)
398-
elif acq_software == "Open Ephys":
399-
dataset = openephys.OpenEphys(session_dir)
400-
for serial_number, probe_data in dataset.probes.items():
401-
if str(serial_number) == inserted_probe_serial_number:
402-
break
403-
else:
404-
raise FileNotFoundError(
405-
"No Open Ephys data found for probe insertion: {}".format(key)
406-
)
407406

408-
if not probe_data.ap_meta:
409-
raise IOError(
410-
'No analog signals found - check "structure.oebin" file or "continuous" directory'
411-
)
407+
if not probe_data.ap_meta:
408+
raise IOError(
409+
'No analog signals found - check "structure.oebin" file or "continuous" directory'
410+
)
412411

413-
if probe_data.probe_model in supported_probe_types:
414-
probe_type = probe_data.probe_model
415-
electrode_query = probe.ProbeType.Electrode & {
416-
"probe_type": probe_type
417-
}
412+
if probe_data.probe_model in supported_probe_types:
413+
probe_type = probe_data.probe_model
414+
electrode_query = probe.ProbeType.Electrode & {
415+
"probe_type": probe_type
416+
}
418417

419-
probe_electrodes = {
420-
key["electrode"]: key for key in electrode_query.fetch("KEY")
421-
}
418+
probe_electrodes = {
419+
key["electrode"]: key for key in electrode_query.fetch("KEY")
420+
}
422421

423-
electrode_group_members = [
424-
probe_electrodes[channel_idx]
425-
for channel_idx in probe_data.ap_meta["channels_indices"]
426-
]
427-
else:
428-
raise NotImplementedError(
429-
"Processing for neuropixels"
430-
" probe model {} not yet implemented".format(
431-
probe_data.probe_model
432-
)
422+
electrode_group_members = [
423+
probe_electrodes[channel_idx]
424+
for channel_idx in probe_data.ap_meta["channels_indices"]
425+
]
426+
else:
427+
raise NotImplementedError(
428+
"Processing for neuropixels"
429+
" probe model {} not yet implemented".format(
430+
probe_data.probe_model
433431
)
434-
435-
self.insert1(
436-
{
437-
**key,
438-
**generate_electrode_config(
439-
probe_type, electrode_group_members
440-
),
441-
"acq_software": acq_software,
442-
"sampling_rate": probe_data.ap_meta["sample_rate"],
443-
"recording_datetime": probe_data.recording_info[
444-
"recording_datetimes"
445-
][0],
446-
"recording_duration": np.sum(
447-
probe_data.recording_info["recording_durations"]
448-
),
449-
}
450432
)
451433

452-
root_dir = find_root_directory(
453-
get_ephys_root_data_dir(),
454-
probe_data.recording_info["recording_files"][0],
455-
)
456-
self.EphysFile.insert(
457-
[
458-
{**key, "file_path": fp.relative_to(root_dir).as_posix()}
459-
for fp in probe_data.recording_info["recording_files"]
460-
]
461-
)
462-
# Explicitly garbage collect "dataset" as these may have large memory footprint and may not be cleared fast enough
463-
del probe_data, dataset
464-
gc.collect()
434+
self.insert1(
435+
{
436+
**key,
437+
**generate_electrode_config(
438+
probe_type, electrode_group_members
439+
),
440+
"acq_software": acq_software,
441+
"sampling_rate": probe_data.ap_meta["sample_rate"],
442+
"recording_datetime": probe_data.recording_info[
443+
"recording_datetimes"
444+
][0],
445+
"recording_duration": np.sum(
446+
probe_data.recording_info["recording_durations"]
447+
),
448+
}
449+
)
465450

466-
# Insert channel information
467-
# Get channel and electrode-site mapping
468-
channel2electrodes = get_neuropixels_channel2electrode_map(
469-
key, acq_software
451+
root_dir = find_root_directory(
452+
get_ephys_root_data_dir(),
453+
probe_data.recording_info["recording_files"][0],
470454
)
471-
self.Channel.insert(
455+
self.EphysFile.insert(
472456
[
473-
{**key, "channel_idx": channel_idx, **channel_info}
474-
for channel_idx, channel_info in channel2electrodes.items()
457+
{**key, "file_path": fp.relative_to(root_dir).as_posix()}
458+
for fp in probe_data.recording_info["recording_files"]
475459
]
476460
)
461+
# Explicitly garbage collect "dataset" as these may have large memory footprint and may not be cleared fast enough
462+
del probe_data, dataset
463+
gc.collect()
464+
465+
# Insert channel information
466+
# Get channel and electrode-site mapping
467+
channel2electrodes = get_neuropixels_channel2electrode_map(
468+
key, acq_software
469+
)
470+
self.Channel.insert(
471+
[
472+
{**key, "channel_idx": channel_idx, **channel_info}
473+
for channel_idx, channel_info in channel2electrodes.items()
474+
]
475+
)
477476

478477

479478
@schema

0 commit comments

Comments
 (0)