Return OSeMOSYS data as an xarray.DataSet

requirements.txt

@@ -8,4 +8,5 @@ pandas
 pydantic
 pydot
 pyyaml
+xarray
 xlrd

setup.cfg

@@ -52,6 +52,7 @@ install_requires =
     flatten_dict
     openpyxl
     pydantic
+    xarray
 [options.packages.find]
 where = src
 exclude =

src/otoole/input.py

@@ -35,6 +35,7 @@
 from typing import Any, Dict, List, Optional, TextIO, Tuple, Union
 
 import pandas as pd
+import xarray as xr
 
 from otoole.exceptions import OtooleIndexError, OtooleNameMismatchError
 
@@ -593,3 +594,34 @@ def read(
         self, filepath: Union[str, TextIO], **kwargs
     ) -> Tuple[Dict[str, pd.DataFrame], Dict[str, Any]]:
         raise NotImplementedError()
+
+    def to_xarray(self, filepath) -> xr.Dataset:
+        """Returns input data as an xarray.Dataset
+
+        Arguments
+        ---------
+        filepath: Union[str, TextIO]
+
+        """
+
+        model, defaults = self.read(filepath)
+        config = self.input_config
+
+        data_vars = {
+            x: y.VALUE.to_xarray()
+            for x, y in model.items()
+            if config[x]["type"] == "param"
+        }
+        coords = {
+            x: y.values.T[0] for x, y in model.items() if config[x]["type"] == "set"
+        }
+        ds = xr.Dataset(data_vars=data_vars, coords=coords)
+        # ds = ds.assign_coords({'_REGION': model['REGION'].values.T[0]})
+
+        for param, default in defaults.items():
+            if param in config and param in model and config[param]["type"] == "param":
+                ds[param].attrs["default"] = default
+                if default != 0:
+                    ds[param] = ds[param].fillna(default)
+
+        return ds

tests/test_input.py

@@ -1,11 +1,13 @@
 from typing import Any, Dict, TextIO, Tuple, Union
 
 import pandas as pd
+import xarray as xr
 from pandas.testing import assert_frame_equal
 from pytest import fixture, mark, raises
 
 from otoole.exceptions import OtooleIndexError, OtooleNameMismatchError
 from otoole.input import ReadStrategy, WriteStrategy
+from otoole.read_strategies import ReadMemory
 
 
 @fixture
@@ -533,3 +535,21 @@ def test_compare_read_to_expected_exception(self, simple_user_config, expected):
         reader = DummyReadStrategy(simple_user_config)
         with raises(OtooleNameMismatchError):
             reader._compare_read_to_expected(names=expected)
+
+
+class TestXarray:
+    def test_xarray(self, user_config):
+        """Test that xarray can be imported"""
+
+        data = [
+            ["SIMPLICITY", "ETH", 2014, 1.0],
+            ["SIMPLICITY", "RAWSUG", 2014, 0.5],
+            ["SIMPLICITY", "ETH", 2015, 1.03],
+            ["SIMPLICITY", "RAWSUG", 2015, 0.51],
+        ]
+        df = pd.DataFrame(data=data, columns=["REGION", "FUEL", "YEAR", "VALUE"])
+        parameters = {"AccumulatedAnnualDemand": df}
+
+        reader = ReadMemory(parameters, user_config)
+        actual = reader.to_xarray("")
+        assert isinstance(actual, xr.Dataset)