88def 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