Explicitly set xarray coordinates.

AllenInstitute · Oct 25, 2023 · 7c748b1 · 7c748b1
1 parent 028a9e6
commit 7c748b1
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/allensdk/brain_observatory/ecephys/stimulus_analysis/receptive_field_mapping.py b/allensdk/brain_observatory/ecephys/stimulus_analysis/receptive_field_mapping.py
@@ -225,8 +225,12 @@ def _response_by_stimulus_position(self, dataset, presentations, row_key=None, c
         dataset[spike_count_key] = dataset.sum(dim=time_key)
         dataset = dataset.drop(time_key)
 
-        dataset[row_key] = presentations.loc[:, row_key]
-        dataset[column_key] = presentations.loc[:, column_key]
+        dataset = dataset.assign_coords(
+            {row_key: ("stimulus_presentation_id",
+                       presentations.loc[:, row_key].to_numpy())})
+        dataset = dataset.assign_coords(
+            {column_key: ("stimulus_presentation_id",
+                          presentations.loc[:, column_key].to_numpy())})
         dataset = dataset.to_dataframe()
 
         dataset = dataset.reset_index(unit_key).groupby([row_key, column_key, unit_key]).sum()