Finish initial market reader

Author: tsmbland

src/muse/data/example/default_new_input/timeslices.csv

@@ -1,7 +1,7 @@
-id,season,day,time_of_day,fraction
-1,all,all,night,0.1667
-2,all,all,morning,0.1667
-3,all,all,afternoon,0.1667
-4,all,all,early-peak,0.1667
-5,all,all,late-peak,0.1667
-6,all,all,evening,0.1667
+id,month,day,hour,fraction
+1,all-year,all-week,night,0.1667
+2,all-year,all-week,morning,0.1667
+3,all-year,all-week,afternoon,0.1667
+4,all-year,all-week,early-peak,0.1667
+5,all-year,all-week,late-peak,0.1667
+6,all-year,all-week,evening,0.1667

src/muse/new_input/readers.py

@@ -3,6 +3,8 @@
 import pandas as pd
 import xarray as xr
 
+from muse.timeslices import QuantityType
+
 
 def read_inputs(data_dir):
     data = {}
@@ -42,17 +44,15 @@ def read_inputs(data_dir):
 def read_timeslices_csv(buffer_, con):
     sql = """CREATE TABLE timeslices (
       id BIGINT PRIMARY KEY,
-      season VARCHAR,
+      month VARCHAR,
       day VARCHAR,
-      time_of_day VARCHAR,
+      hour VARCHAR,
       fraction DOUBLE CHECK (fraction >= 0 AND fraction <= 1),
     );
     """
     con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
-    con.sql(
-        "INSERT INTO timeslices SELECT id, season, day, time_of_day, fraction FROM rel;"
-    )
+    con.sql("INSERT INTO timeslices SELECT id, month, day, hour, fraction FROM rel;")
     return con.sql("SELECT * from timeslices").fetchnumpy()
 
 
@@ -278,17 +278,19 @@ def calculate_initial_market(
     - If price data is not specified for a commodity/region combination, the price is
     zero
 
-    """
-    from muse.timeslices import QuantityType, convert_timeslice
+    Todo:
+    - Allow data to be specified on a timeslice level (optional)
+    - Interpolation, missing year field, flexible timeslice specification as above
 
+    """
     # Prepare dataframes
     df_trade = pd.DataFrame(commodity_trade).set_index(["commodity", "region", "year"])
     df_costs = (
         pd.DataFrame(commodity_costs)
         .set_index(["commodity", "region", "year"])
         .rename(columns={"value": "prices"})
     )
-    df_timeslices = pd.DataFrame(timeslices).set_index(["season", "day", "time_of_day"])
+    df_timeslices = pd.DataFrame(timeslices).set_index(["month", "day", "hour"])
 
     # DataArray dimensions
     all_commodities = commodities["id"].astype(np.dtype("str"))
@@ -320,13 +322,17 @@ def calculate_initial_market(
     # Calculate static trade
     df_trade["static_trade"] = df_trade["export"] - df_trade["import"]
 
-    # Create Data
-    df_full = df_costs.join(df_trade)
-    data = df_full.to_xarray()
-    ts = df_timeslices.to_xarray()["fraction"]
-    ts = ts.stack(timeslice=("season", "day", "time_of_day"))
-    convert_timeslice(data, ts, QuantityType.EXTENSIVE)
+    # Create xarray datasets
+    xr_costs = df_costs.to_xarray()
+    xr_trade = df_trade.to_xarray()
+
+    # Project over timeslices
+    ts = df_timeslices.to_xarray()["fraction"].stack(timeslice=("month", "day", "hour"))
+    xr_costs = project_timeslice(xr_costs, ts, QuantityType.EXTENSIVE)
+    xr_trade = project_timeslice(xr_trade, ts, QuantityType.INTENSIVE)
 
+    # Combine data
+    data = xr.merge([xr_costs, xr_trade])
     return data
 
 
@@ -353,3 +359,30 @@ def check_all_values_specified(
     ).all():
         msg = ""  # TODO
         raise DataValidationError(msg)
+
+
+def project_timeslice(
+    data: xr.Dataset, timeslices: xr.DataArray, quantity_type: QuantityType
+) -> xr.Dataset:
+    """Project a dataset over a new timeslice dimension.
+
+    The projection can be done in one of two ways, depending on whether the
+    quantity type is extensive or intensive. See `QuantityType`.
+
+    Args:
+        data: Dataset to project
+        timeslices: DataArray of timeslice levels, with values between 0 and 1
+            representing the timeslice length (fraction of the year)
+        quantity_type: Type of projection to perform. QuantityType.EXTENSIVE or
+            QuantityType.INTENSIVE
+
+    Returns:
+        Projected dataset
+    """
+    assert "timeslice" in timeslices.dims
+    assert "timeslice" not in data.dims
+
+    if quantity_type is QuantityType.INTENSIVE:
+        return data * timeslices
+    if quantity_type is QuantityType.EXTENSIVE:
+        return data * xr.ones_like(timeslices)

tests/test_new_readers.py

@@ -3,7 +3,7 @@
 import duckdb
 import numpy as np
 import xarray as xr
-from pytest import approx, fixture, mark, raises
+from pytest import approx, fixture, raises
 
 
 @fixture
@@ -206,9 +206,9 @@ def test_calculate_demand(
     assert data.dtype == np.float64
 
     assert set(data.dims) == {"year", "commodity", "region", "timeslice"}
-    assert list(data.coords["region"].values) == ["R1"]
+    assert set(data.coords["region"].values) == {"R1"}
     assert set(data.coords["timeslice"].values) == set(range(1, 7))
-    assert list(data.coords["year"].values) == [2020, 2050]
+    assert set(data.coords["year"].values) == {2020, 2050}
     assert set(data.coords["commodity"].values) == {
         "electricity",
         "gas",
@@ -220,7 +220,6 @@ def test_calculate_demand(
     assert data.sel(year=2020, commodity="heat", region="R1", timeslice=1) == 1
 
 
-@mark.xfail
 def test_calculate_initial_market(
     populate_commodities,
     populate_regions,
@@ -240,12 +239,8 @@ def test_calculate_initial_market(
 
     assert isinstance(data, xr.Dataset)
     assert set(data.dims) == {"region", "year", "commodity", "timeslice"}
-    assert dict(data.dtypes) == dict(
-        prices=np.float64,
-        exports=np.float64,
-        imports=np.float64,
-        static_trade=np.float64,
-    )
+    for dt in data.dtypes.values():
+        assert dt == np.dtype("float64")
     assert set(data.coords["region"].values) == {"R1"}
     assert set(data.coords["year"].values) == set(range(2010, 2105, 5))
     assert set(data.coords["commodity"].values) == {
@@ -266,28 +261,30 @@ def test_calculate_initial_market(
         "evening",
     ]
 
-    assert list(data.coords["timeslice"].values) == list(
+    assert set(data.coords["timeslice"].values) == set(
         zip(month_values, day_values, hour_values)
     )
-    assert list(data.coords["month"]) == month_values
-    assert list(data.coords["day"]) == day_values
-    assert list(data.coords["hour"]) == hour_values
+    assert set(data.coords["month"].values) == set(month_values)
+    assert set(data.coords["day"].values) == set(day_values)
+    assert set(data.coords["hour"].values) == set(hour_values)
 
     assert all(var.coords.equals(data.coords) for var in data.data_vars.values())
 
     prices = data.data_vars["prices"]
-    assert approx(
-        prices.sel(
-            year=2010,
-            region="R1",
-            commodity="electricity",
-            timeslice=("all-year", "all-week", "night"),
+    assert (
+        approx(
+            prices.sel(
+                year=2010,
+                region="R1",
+                commodity="electricity",
+                timeslice=("all-year", "all-week", "night"),
+            ),
+            abs=1e-4,
         )
-        - 14.81481,
-        abs=1e-4,
+        == 14.81481
     )
 
-    exports = data.data_vars["exports"]
+    exports = data.data_vars["export"]
     assert (
         exports.sel(
             year=2010,
@@ -297,7 +294,7 @@ def test_calculate_initial_market(
         )
     ) == 0
 
-    imports = data.data_vars["imports"]
+    imports = data.data_vars["import"]
     assert (
         imports.sel(
             year=2010,