Merge pull request #189 from neurolib-dev/fix/multimodel_exploration

Quick fixes for BoxSearch xr() when exploring MultiModel

Author: caglorithm

neurolib/optimize/exploration/exploration.py

@@ -488,6 +488,10 @@ def xr(self, bold=False):
         :param bold: if True, will load and return only BOLD output
         :type bold: bool
         """
+
+        def _sanitize_nc_key(k):
+            return k.replace("*", "_").replace(".", "_").replace("|", "_")
+
         assert self.results is not None, "Run `loadResults()` first to populate the results"
         assert len(self.results) == len(self.dfResults)
         # create intrisinsic dims for one run
@@ -510,6 +514,8 @@ def xr(self, bold=False):
             expand_coords = {}
             # iterate exploration coordinates
             for k, v in expl_coords.items():
+                # sanitize keys in the case of stars etc
+                k = _sanitize_nc_key(k)
                 # if single values, just assing
                 if isinstance(v, (str, float, int)):
                     expand_coords[k] = [v]
@@ -525,10 +531,14 @@ def xr(self, bold=False):
             dataarrays.append(data_temp.expand_dims(expand_coords))
 
         # finally, combine all arrays into one
-        combined = xr.combine_by_coords(dataarrays)["exploration"]
+        # sometimes combining xr.DataArrays does not work, see https://github.com/pydata/xarray/issues/3248#issuecomment-531511177
+        # resolved by casting them explicitely to xr.Dataset
+        combined = xr.combine_by_coords([da.to_dataset() for da in dataarrays])["exploration"]
         if self.parameterSpace.star:
-            combined.attrs = {k: list(self.model.params[k].keys()) for k in orig_search_coords.keys()}
-
+            # if we explored over star params, unwrap them into attributes
+            combined.attrs = {
+                _sanitize_nc_key(k): list(self.model.params[k].keys()) for k in orig_search_coords.keys() if "*" in k
+            }
         return combined
 
     def getRun(self, runId, filename=None, trajectoryName=None, pypetShortNames=True):

tests/test_exploration.py

@@ -188,7 +188,10 @@ def test_multimodel_explore(self):
         dataarray = search.xr()
         self.assertTrue(isinstance(dataarray, xr.DataArray))
         self.assertTrue(isinstance(dataarray.attrs, dict))
-        self.assertListEqual(list(dataarray.attrs.keys()), list(parameters.dict().keys()))
+        self.assertListEqual(
+            list(dataarray.attrs.keys()),
+            [k.replace("*", "_").replace(".", "_").replace("|", "_") for k in parameters.dict().keys()],
+        )
 
         end = time.time()
         logging.info("\t > Done in {:.2f} s".format(end - start))