|
| 1 | +import importlib.util |
| 2 | +import os |
| 3 | +import numpy as np |
| 4 | + |
| 5 | +from ..basesorter import BaseSorter |
| 6 | + |
| 7 | +from spikeinterface.extractors import NumpySorting # TODO: Create separate sorting extractor for RT-Sort |
| 8 | + |
| 9 | + |
| 10 | +class RTSortSorter(BaseSorter): |
| 11 | + """RTSort sorter object""" |
| 12 | + |
| 13 | + sorter_name = "rt-sort" |
| 14 | + |
| 15 | + _default_params = { |
| 16 | + "detection_model": "neuropixels", |
| 17 | + "recording_window_ms": None, |
| 18 | + "stringent_thresh": 0.175, |
| 19 | + "loose_thresh": 0.075, |
| 20 | + "inference_scaling_numerator": 15.4, |
| 21 | + "ms_before": 0.5, |
| 22 | + "ms_after": 0.5, |
| 23 | + "pre_median_ms": 50, |
| 24 | + "inner_radius": 50, |
| 25 | + "outer_radius": 100, |
| 26 | + "min_elecs_for_array_noise_n": 100, |
| 27 | + "min_elecs_for_array_noise_f": 0.1, |
| 28 | + "min_elecs_for_seq_noise_n": 50, |
| 29 | + "min_elecs_for_seq_noise_f": 0.05, |
| 30 | + "min_activity_root_cocs": 2, |
| 31 | + "min_activity_hz": 0.05, |
| 32 | + "max_n_components_latency": 4, |
| 33 | + "min_coc_n": 10, |
| 34 | + "min_coc_p": 10, |
| 35 | + "min_extend_comp_p": 50, |
| 36 | + "elec_patience": 6, |
| 37 | + "split_coc_clusters_amps": True, |
| 38 | + "min_amp_dist_p": 0.1, |
| 39 | + "max_n_components_amp": 4, |
| 40 | + "min_loose_elec_prob": 0.03, |
| 41 | + "min_inner_loose_detections": 3, |
| 42 | + "min_loose_detections_n": 4, |
| 43 | + "min_loose_detections_r_spikes": 1 / 3, |
| 44 | + "min_loose_detections_r_sequences": 1 / 3, |
| 45 | + "max_latency_diff_spikes": 2.5, |
| 46 | + "max_latency_diff_sequences": 2.5, |
| 47 | + "clip_latency_diff_factor": 2, |
| 48 | + "max_amp_median_diff_spikes": 0.45, |
| 49 | + "max_amp_median_diff_sequences": 0.45, |
| 50 | + "clip_amp_median_diff_factor": 2, |
| 51 | + "max_root_amp_median_std_spikes": 2.5, |
| 52 | + "max_root_amp_median_std_sequences": 2.5, |
| 53 | + "repeated_detection_overlap_time": 0.2, |
| 54 | + "min_seq_spikes_n": 10, |
| 55 | + "min_seq_spikes_hz": 0.05, |
| 56 | + "relocate_root_min_amp": 0.8, |
| 57 | + "relocate_root_max_latency": -2, |
| 58 | + "device": "cuda", |
| 59 | + "num_processes": None, |
| 60 | + "ignore_warnings": True, |
| 61 | + "debug": False, |
| 62 | + } |
| 63 | + |
| 64 | + _params_description = { |
| 65 | + "detection_model": "`mea` or `neuropixels` to use the mea- or neuropixels-trained detection models. Or, path to saved detection model (see https://braingeneers.github.io/braindance/docs/RT-sort/usage/training-models) (`str`, `neuropixels`)", |
| 66 | + "recording_window_ms": "A tuple `(start_ms, end_ms)` defining the portion of the recording (in milliseconds) to process. (`tuple`, `None`)", |
| 67 | + "stringent_thresh": "The stringent threshold for spike detection. (`float`, `0.175`)", |
| 68 | + "loose_thresh": "The loose threshold for spike detection. (`float`, `0.075`)", |
| 69 | + "inference_scaling_numerator": "Scaling factor for inference. (`float`, `15.4`)", |
| 70 | + "ms_before": "Time (in milliseconds) to consider before each detected spike for sequence formation. (`float`, `0.5`)", |
| 71 | + "ms_after": "Time (in milliseconds) to consider after each detected spike for sequence formation. (`float`, `0.5`)", |
| 72 | + "pre_median_ms": "Duration (in milliseconds) to compute the median for normalization. (`float`, `50`)", |
| 73 | + "inner_radius": "Inner radius (in micrometers). (`float`, `50`)", |
| 74 | + "outer_radius": "Outer radius (in micrometers). (`float`, `100`)", |
| 75 | + "min_elecs_for_array_noise_n": "Minimum number of electrodes for array-wide noise filtering. (`int`, `100`)", |
| 76 | + "min_elecs_for_array_noise_f": "Minimum fraction of electrodes for array-wide noise filtering. (`float`, `0.1`)", |
| 77 | + "min_elecs_for_seq_noise_n": "Minimum number of electrodes for sequence-wide noise filtering. (`int`, `50`)", |
| 78 | + "min_elecs_for_seq_noise_f": "Minimum fraction of electrodes for sequence-wide noise filtering. (`float`, `0.05`)", |
| 79 | + "min_activity_root_cocs": "Minimum number of stringent spike detections on inner electrodes within the maximum propagation window that cause a stringent spike detection on a root electrode to be counted as a stringent codetection. (`int`, `2`)", |
| 80 | + "min_activity_hz": "Minimum activity rate of root detections (in Hz) for an electrode to be used as a root electrode. (`float`, `0.05`)", |
| 81 | + "max_n_components_latency": "Maximum number of latency components for Gaussian mixture model used for splitting latency distribution. (`int`, `4`)", |
| 82 | + "min_coc_n": "After splitting a cluster of codetections, a cluster is discarded if it does not have at least min_coc_n codetections. (`int`, `10`)", |
| 83 | + "min_coc_p": "After splitting a cluster of codetections, a cluster is discarded if it does not have at least (min_coc_p * the total number of codetections before splitting) codetections. (`int`, `10`)", |
| 84 | + "min_extend_comp_p": "The required percentage of codetections before splitting that is preserved after the split in order for the inner electrodes of the current splitting electrode to be added to the total list of electrodes used to further split the cluster. (`int`, `50`)", |
| 85 | + "elec_patience": "Number of electrodes considered for splitting that do not lead to a split before terminating the splitting process. (`int`, `6`)", |
| 86 | + "split_coc_clusters_amps": "Whether to split clusters based on amplitude. (`bool`, `True`)", |
| 87 | + "min_amp_dist_p": "The minimum Hartigan's dip test p-value for a distribution to be considered unimodal. (`float`, `0.1`)", |
| 88 | + "max_n_components_amp": "Maximum number of components for Gaussian mixture model used for splitting amplitude distribution. (`int`, `4`)", |
| 89 | + "min_loose_elec_prob": "Minimum average detection score (smaller values are set to 0) in decimal form (ranging from 0 to 1). (`float`, `0.03`)", |
| 90 | + "min_inner_loose_detections": "Minimum inner loose electrode detections for assigning spikes / overlaps for merging. (`int`, `3`)", |
| 91 | + "min_loose_detections_n": "Minimum loose electrode detections for assigning spikes / overlaps for merging. (`int`, `4`)", |
| 92 | + "min_loose_detections_r_spikes": "Minimum ratio of loose electrode detections for assigning spikes. (`float`, `1/3`)", |
| 93 | + "min_loose_detections_r_sequences": "Minimum ratio of loose electrode detections overlaps for merging. (`float`, `1/3`)", |
| 94 | + "max_latency_diff_spikes": "Maximum allowed weighted latency difference for spike assignment. (`float`, `2.5`)", |
| 95 | + "max_latency_diff_sequences": "Maximum allowed weighted latency difference for sequence merging. (`float`, `2.5`)", |
| 96 | + "clip_latency_diff_factor": "Latency clip = clip_latency_diff_factor * max_latency_diff. (`float`, `2`)", |
| 97 | + "max_amp_median_diff_spikes": "Maximum allowed weighted percent amplitude difference for spike assignment. (`float`, `0.45`)", |
| 98 | + "max_amp_median_diff_sequences": "Maximum allowed weighted percent amplitude difference for sequence merging. (`float`, `0.45`)", |
| 99 | + "clip_amp_median_diff_factor": "Amplitude clip = clip_amp_median_diff_factor * max_amp_median_diff. (`float`, `2`)", |
| 100 | + "max_root_amp_median_std_spikes": "Maximum allowed root amplitude standard deviation for spike assignment. (`float`, `2.5`)", |
| 101 | + "max_root_amp_median_std_sequences": "Maximum allowed root amplitude standard deviation for sequence merging. (`float`, `2.5`)", |
| 102 | + "repeated_detection_overlap_time": "Time window (in seconds) for overlapping repeated detections. (`float`, `0.2`)", |
| 103 | + "min_seq_spikes_n": "Minimum number of spikes required for a valid sequence. (`int`, `10`)", |
| 104 | + "min_seq_spikes_hz": "Minimum spike rate for a valid sequence. (`float`, `0.05`)", |
| 105 | + "relocate_root_min_amp": "Minimum amplitude ratio for relocating a root electrode before first merging. (`float`, `0.8`)", |
| 106 | + "relocate_root_max_latency": "Maximum latency for relocating a root electrode before first merging. (`float`, `-2`)", |
| 107 | + "device": "The device for PyTorch operations ('cuda' or 'cpu'). (`str`, `cuda`)", |
| 108 | + "num_processes": "Number of processes to use for parallelization. (`int`, `None`)", |
| 109 | + "ignore_warnings": "Whether to suppress warnings during execution. (`bool`, `True`)", |
| 110 | + "debug": "Whether to enable debugging features such as saving intermediate steps. (`bool`, `False`)", |
| 111 | + } |
| 112 | + |
| 113 | + sorter_description = """RT-Sort is a real-time spike sorting algorithm that enables the sorted detection of action potentials within 7.5ms±1.5ms (mean±STD) after the waveform trough while the recording remains ongoing. |
| 114 | + It utilizes unique propagation patterns of action potentials along axons detected as high-fidelity sequential activations on adjacent electrodes, together with a convolutional neural network-based spike detection algorithm. |
| 115 | + This implementation in SpikeInterface only implements RT-Sort's offline sorting. |
| 116 | + For more information see https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0312438""" |
| 117 | + |
| 118 | + installation_mesg = f"""\nTo use RTSort run:\n |
| 119 | + >>> pip install git+https://github.com/braingeneers/braindance#egg=braindance[rt-sort] |
| 120 | +
|
| 121 | + Additionally, install PyTorch (https://pytorch.org/get-started/locally/) with any version of CUDA as the compute platform. |
| 122 | + If running on a Linux machine, install Torch-TensorRT (https://pytorch.org/TensorRT/getting_started/installation.html) for faster computations. |
| 123 | +
|
| 124 | + More information on RTSort at: https://github.com/braingeneers/braindance |
| 125 | + """ |
| 126 | + |
| 127 | + handle_multi_segment = False |
| 128 | + |
| 129 | + @classmethod |
| 130 | + def get_sorter_version(cls): |
| 131 | + import braindance |
| 132 | + |
| 133 | + return braindance.__version__ |
| 134 | + |
| 135 | + @classmethod |
| 136 | + def is_installed(cls): |
| 137 | + libraries = ["braindance", "torch" if os.name == "nt" else "torch_tensorrt", "diptest", "pynvml", "sklearn"] |
| 138 | + |
| 139 | + HAVE_RTSORT = True |
| 140 | + for lib in libraries: |
| 141 | + if importlib.util.find_spec(lib) is None: |
| 142 | + HAVE_RTSORT = False |
| 143 | + break |
| 144 | + |
| 145 | + return HAVE_RTSORT |
| 146 | + |
| 147 | + @classmethod |
| 148 | + def _check_params(cls, recording, output_folder, params): |
| 149 | + return params |
| 150 | + |
| 151 | + @classmethod |
| 152 | + def _check_apply_filter_in_params(cls, params): |
| 153 | + return False |
| 154 | + |
| 155 | + @classmethod |
| 156 | + def _setup_recording(cls, recording, sorter_output_folder, params, verbose): |
| 157 | + # nothing to copy inside the folder : RTSort uses spikeinterface natively |
| 158 | + pass |
| 159 | + |
| 160 | + @classmethod |
| 161 | + def _run_from_folder(cls, sorter_output_folder, params, verbose): |
| 162 | + from braindance.core.spikesorter.rt_sort import detect_sequences |
| 163 | + from braindance.core.spikedetector.model import ModelSpikeSorter |
| 164 | + |
| 165 | + recording = cls.load_recording_from_folder(sorter_output_folder.parent, with_warnings=False) |
| 166 | + |
| 167 | + params = params.copy() |
| 168 | + params["recording"] = recording |
| 169 | + rt_sort_inter = sorter_output_folder / "rt_sort_inter" |
| 170 | + params["inter_path"] = rt_sort_inter |
| 171 | + params["verbose"] = verbose |
| 172 | + |
| 173 | + if params["detection_model"] == "mea": |
| 174 | + params["detection_model"] = ModelSpikeSorter.load_mea() |
| 175 | + elif params["detection_model"] == "neuropixels": |
| 176 | + params["detection_model"] = ModelSpikeSorter.load_neuropixels() |
| 177 | + else: |
| 178 | + params["detection_model"] = ModelSpikeSorter.load(params["detection_model"]) |
| 179 | + |
| 180 | + rt_sort = detect_sequences(**params, delete_inter=False, return_spikes=False) |
| 181 | + np_sorting = rt_sort.sort_offline( |
| 182 | + rt_sort_inter / "scaled_traces.npy", |
| 183 | + verbose=verbose, |
| 184 | + recording_window_ms=params.get("recording_window_ms", None), |
| 185 | + return_spikeinterface_sorter=True, |
| 186 | + ) # type: NumpySorting |
| 187 | + rt_sort.save(sorter_output_folder / "rt_sort.pickle") |
| 188 | + np_sorting.save(folder=sorter_output_folder / "rt_sorting") |
| 189 | + |
| 190 | + @classmethod |
| 191 | + def _get_result_from_folder(cls, sorter_output_folder): |
| 192 | + return NumpySorting.load_from_folder(sorter_output_folder / "rt_sorting") |
0 commit comments