Skip to content

Commit dd063c6

Browse files
to_pynapple_tsgroup : attach_unit_properties, export sorting properties (#4360)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 3eb9381 commit dd063c6

1 file changed

Lines changed: 15 additions & 3 deletions

File tree

src/spikeinterface/exporters/to_pynapple.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
def to_pynapple_tsgroup(
99
sorting_analyzer_or_sorting: SortingAnalyzer | BaseSorting,
1010
attach_unit_metadata=True,
11+
attach_unit_properties=True,
1112
segment_index=None,
1213
):
1314
"""
@@ -21,6 +22,8 @@ def to_pynapple_tsgroup(
2122
If True, any relevant available metadata is attached to the TsGroup. Will attach
2223
`unit_locations`, `quality_metrics` and `template_metrics` if computed. If False,
2324
no metadata is included.
25+
attach_unit_properties : bool, default: False
26+
If True, attach properties of the sorting.
2427
segment_index : int | None, default: None
2528
The segment index. Can be None if mono-segment sorting.
2629
@@ -59,14 +62,13 @@ def to_pynapple_tsgroup(
5962
for unit_id_int, unit_id in zip(unit_ids_ints, unit_ids)
6063
}
6164

62-
metadata_list = []
65+
metadata_list = [] # init list to collect metadata dataframes
66+
6367
if not unit_ids_castable:
6468
metadata_list.append(pd.DataFrame(unit_ids, columns=["unit_id"]))
6569

6670
# Look for good metadata to add, if there is a sorting analyzer
6771
if attach_unit_metadata and isinstance(sorting_analyzer_or_sorting, SortingAnalyzer):
68-
69-
metadata_list = []
7072
if (unit_locations := sorting_analyzer_or_sorting.get_extension("unit_locations")) is not None:
7173
array_of_unit_locations = unit_locations.get_data()
7274
n_dims = np.shape(sorting_analyzer_or_sorting.get_extension("unit_locations").get_data())[1]
@@ -79,6 +81,16 @@ def to_pynapple_tsgroup(
7981
if (template_metrics := sorting_analyzer_or_sorting.get_extension("template_metrics")) is not None:
8082
metadata_list.append(template_metrics.get_data())
8183

84+
# attach unit properties from sorting
85+
if attach_unit_properties:
86+
property_df = pd.DataFrame(index=unit_ids)
87+
property_keys = sorting.get_property_keys() # get property keys of sorting
88+
if len(property_keys):
89+
for property_key in property_keys: # loop through sorting's properties
90+
property_data = sorting.get_property(property_key)
91+
property_df[property_key] = list(property_data)
92+
metadata_list.append(property_df)
93+
8294
if len(metadata_list) > 0:
8395
metadata = pd.concat(metadata_list, axis=1)
8496
metadata.index = unit_ids_ints

0 commit comments

Comments
 (0)