# HG changeset patch
# User Ashwin Vishnu <avmo@kth.se>
# Date 1541368610 -3600
#      Sun Nov 04 22:56:50 2018 +0100
# Node ID 5118e5a668a8c36e767e81ef549c5d82b73ebb70
# Parent  19b7531f78b8763e976035e7f3b612283996e17d
Load spatial means as an xarray dataset

diff --git a/fluidsim/base/output/spatial_means.py b/fluidsim/base/output/spatial_means.py
--- a/fluidsim/base/output/spatial_means.py
+++ b/fluidsim/base/output/spatial_means.py
@@ -103,6 +103,23 @@
         dict_results = {}
         return dict_results
 
+    def load_dataset(self, dims=("t",)):
+        """Loads results as a xarray dataset."""
+        dict_results = self.load()
+        # NOTE: format specified in
+        # http://xarray.pydata.org/en/stable/generated/xarray.Dataset.from_dict.html
+        dset = {"coords": {}, "attrs": {}, "dims": dims, "data_vars": {}}
+        for key, value in dict_results.items():
+            if isinstance(value, np.ndarray):
+                target = "coords" if key in dims else "data_vars"
+                dset[target].update({key: {"dims": dims, "data": value}})
+            else:
+                dset["attrs"].update({key: value})
+
+        from xarray import Dataset
+
+        return Dataset.from_dict(dset)
+
     def plot(self):
         pass