@@ -191,7 +191,7 @@ def __init__(
191191 self ._dtype = spike_peak_dtype
192192
193193 self .include_spikes_in_margin = include_spikes_in_margin
194- if include_spikes_in_margin is not None :
194+ if include_spikes_in_margin :
195195 self ._dtype = spike_peak_dtype + [("in_margin" , "bool" )]
196196
197197 self .peaks = sorting_to_peaks (sorting , extremum_channel_inds , self ._dtype )
@@ -228,12 +228,6 @@ def compute(self, traces, start_frame, end_frame, segment_index, max_margin, pea
228228 # get local peaks
229229 sl = self .segment_slices [segment_index ]
230230 peaks_in_segment = self .peaks [sl ]
231- # if self.include_spikes_in_margin:
232- # i0, i1 = np.searchsorted(
233- # peaks_in_segment["sample_index"], [start_frame - max_margin, end_frame + max_margin]
234- # )
235- # else:
236- # i0, i1 = np.searchsorted(peaks_in_segment["sample_index"], [start_frame, end_frame])
237231 i0 , i1 = peak_slice
238232
239233 local_peaks = peaks_in_segment [i0 :i1 ]
@@ -435,21 +429,59 @@ def compute(self, traces, peaks):
435429 return sparse_wfs
436430
437431
438- def find_parent_of_type (list_of_parents , parent_type , unique = True ):
432+ def find_parent_of_type (list_of_parents , parent_type ):
433+ """
434+ Find a single parent of a given type(s) in a list of parents.
435+ If multiple parents of the given type are found, the first parent is returned.
436+
437+ Parameters
438+ ----------
439+ list_of_parents : list of PipelineNode
440+ List of parents to search through.
441+ parent_type : type | tuple of types
442+ The type of parent to search for.
443+
444+ Returns
445+ -------
446+ parent : PipelineNode or None
447+ The parent of the given type. Returns None if no parent of the given type is found.
448+ """
439449 if list_of_parents is None :
440450 return None
441451
452+ parents = find_parents_of_type (list_of_parents , parent_type )
453+
454+ if len (parents ) > 0 :
455+ return parents [0 ]
456+ else :
457+ return None
458+
459+
460+ def find_parents_of_type (list_of_parents , parent_type ):
461+ """
462+ Find all parents of a given type(s) in a list of parents.
463+
464+ Parameters
465+ ----------
466+ list_of_parents : list of PipelineNode
467+ List of parents to search through.
468+ parent_type : type | tuple of types
469+ The type(s) of parents to search for.
470+
471+ Returns
472+ -------
473+ parents : list of PipelineNode
474+ List of parents of the given type(s). Returns an empty list if no parents of the given type(s) are found.
475+ """
476+ if list_of_parents is None :
477+ return []
478+
442479 parents = []
443480 for parent in list_of_parents :
444481 if isinstance (parent , parent_type ):
445482 parents .append (parent )
446483
447- if unique and len (parents ) == 1 :
448- return parents [0 ]
449- elif not unique and len (parents ) > 1 :
450- return parents [0 ]
451- else :
452- return None
484+ return parents
453485
454486
455487def check_graph (nodes ):
@@ -471,7 +503,7 @@ def check_graph(nodes):
471503 assert parent in nodes , f"Node { node } has parent { parent } that was not passed in nodes"
472504 assert (
473505 nodes .index (parent ) < i
474- ), f"Node are ordered incorrectly: { node } before { parent } in the pipeline definition."
506+ ), f"Nodes are ordered incorrectly: { node } before { parent } in the pipeline definition."
475507
476508 return nodes
477509
@@ -607,12 +639,16 @@ def _compute_peak_pipeline_chunk(segment_index, start_frame, end_frame, worker_c
607639 skip_after_n_peaks_per_worker = worker_ctx ["skip_after_n_peaks_per_worker" ]
608640
609641 recording_segment = recording ._recording_segments [segment_index ]
610- node0 = nodes [0 ]
611-
612- if isinstance (node0 , (SpikeRetriever , PeakRetriever )):
613- # in this case PeakSource could have no peaks and so no need to load traces just skip
614- peak_slice = i0 , i1 = node0 .get_peak_slice (segment_index , start_frame , end_frame , max_margin )
615- load_trace_and_compute = i0 < i1
642+ retrievers = find_parents_of_type (nodes , (SpikeRetriever , PeakRetriever ))
643+ # get peak slices once for all retrievers
644+ peak_slice_by_retriever = {}
645+ for retriever in retrievers :
646+ peak_slice = i0 , i1 = retriever .get_peak_slice (segment_index , start_frame , end_frame , max_margin )
647+ peak_slice_by_retriever [retriever ] = peak_slice
648+
649+ if len (peak_slice_by_retriever ) > 0 :
650+ # in this case the retrievers could have no peaks, so we test if any spikes are in the chunk
651+ load_trace_and_compute = any (i0 < i1 for i0 , i1 in peak_slice_by_retriever .values ())
616652 else :
617653 # PeakDetector always need traces
618654 load_trace_and_compute = True
@@ -646,7 +682,8 @@ def _compute_peak_pipeline_chunk(segment_index, start_frame, end_frame, worker_c
646682 node_output = node .compute (trace_detection , start_frame , end_frame , segment_index , max_margin )
647683 # set sample index to local
648684 node_output [0 ]["sample_index" ] += extra_margin
649- elif isinstance (node , PeakSource ):
685+ elif isinstance (node , (PeakRetriever , SpikeRetriever )):
686+ peak_slice = peak_slice_by_retriever [node ]
650687 node_output = node .compute (traces_chunk , start_frame , end_frame , segment_index , max_margin , peak_slice )
651688 else :
652689 # TODO later when in master: change the signature of all nodes (or maybe not!)
0 commit comments