Merge pull request #3 from ttngu207/nwb-export

bugfix in assigning unit electrode indices

Author: Thinh Nguyen

element_array_ephys/export/nwb.py

@@ -98,6 +98,9 @@ def curated_clusterings_to_nwb(curated_clustering_keys, nwbfile, ephys_module=No
                 location=electrode_group.location)
 
         # ---- Units ----
+        electrode_df = nwbfile.electrodes.to_dataframe()
+        electrode_ind = electrode_df.index[electrode_df.group_name == electrode_group.name]
+
         unit_query = clustering_query @ ephys.CuratedClustering.Unit
         for unit in unit_query.fetch(as_dict=True):
             waveform_mean = waveform_std = np.full(1, np.nan)
@@ -113,7 +116,7 @@ def curated_clusterings_to_nwb(curated_clustering_keys, nwbfile, ephys_module=No
                     waveform_std = np.std(electrode_waveforms, axis=0)
 
             nwbfile.add_unit(id=unit['unit'],
-                             electrodes=np.where(nwbfile.electrodes.id.data == int(unit['electrode']))[0],
+                             electrodes=np.where(electrode_ind == unit['electrode'])[0],
                              electrode_group=electrode_group,
                              sampling_rate=unit['sampling_rate'],
                              cluster_quality_label=unit['cluster_quality_label'],