First pass at duckdb data interface

Author: cc-a

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

@@ -1,6 +1,6 @@
-id,description,type,unit
+name,description,type,unit
 electricity,Electricity,energy,PJ
 gas,Gas,energy,PJ
 heat,Heat,energy,PJ
 wind,Wind,energy,PJ
-C02f,Carbon dioxide,energy,kt
+CO2f,Carbon dioxide,energy,kt

src/muse/new_input/readers.py

@@ -0,0 +1,76 @@
+import duckdb
+import numpy as np
+import xarray as xr
+
+
+def read_inputs(data_dir):
+    data = {}
+    con = duckdb.connect(":memory:")
+
+    with open(data_dir / "regions.csv") as f:
+        regions = read_regions_csv(f, con)  # noqa: F841
+
+    with open(data_dir / "commodities.csv") as f:
+        commodities = read_commodities_csv(f, con)
+
+    with open(data_dir / "demand.csv") as f:
+        demand = read_demand_csv(f, con)  # noqa: F841
+
+    data["global_commodities"] = calculate_global_commodities(commodities)
+    return data
+
+
+def read_regions_csv(buffer_, con):
+    sql = """CREATE TABLE regions (
+      name VARCHAR PRIMARY KEY,
+    );
+    """
+    con.sql(sql)
+    rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
+    con.sql("INSERT INTO regions SELECT name FROM rel;")
+    return con.sql("SELECT name from regions").fetchnumpy()
+
+
+def read_commodities_csv(buffer_, con):
+    sql = """CREATE TABLE commodities (
+      name VARCHAR PRIMARY KEY,
+      type VARCHAR CHECK (type IN ('energy', 'service', 'material', 'environmental')),
+      unit VARCHAR,
+    );
+    """
+    con.sql(sql)
+    rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
+    con.sql("INSERT INTO commodities SELECT name, type, unit FROM rel;")
+
+    return con.sql("select name, type, unit from commodities").fetchnumpy()
+
+
+def calculate_global_commodities(commodities):
+    names = commodities["name"].astype(np.dtype("str"))
+    types = commodities["type"].astype(np.dtype("str"))
+    units = commodities["unit"].astype(np.dtype("str"))
+
+    type_array = xr.DataArray(
+        data=types, dims=["commodity"], coords=dict(commodity=names)
+    )
+
+    unit_array = xr.DataArray(
+        data=units, dims=["commodity"], coords=dict(commodity=names)
+    )
+
+    data = xr.Dataset(data_vars=dict(type=type_array, unit=unit_array))
+    return data
+
+
+def read_demand_csv(buffer_, con):
+    sql = """CREATE TABLE demand (
+    year BIGINT,
+    commodity VARCHAR REFERENCES commodities(name),
+    region VARCHAR REFERENCES regions(name),
+    demand DOUBLE,
+    );
+    """
+    con.sql(sql)
+    rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
+    con.sql("INSERT INTO demand SELECT year, commodity_name, region, demand FROM rel;")
+    return con.sql("SELECT * from demand").fetchnumpy()

tests/test_readers.py

@@ -1,7 +1,10 @@
+from io import StringIO
 from itertools import chain, permutations
 from pathlib import Path
 
 import pandas as pd
+import duckdb
+import numpy as np
 import toml
 import xarray as xr
 from pytest import fixture, mark, raises
@@ -314,3 +317,147 @@ def test_get_nan_coordinates():
     dataset3 = xr.Dataset.from_dataframe(df3.set_index(["region", "year"]))
     nan_coords3 = get_nan_coordinates(dataset3)
     assert nan_coords3 == []
+
+
+@fixture
+def default_new_input(tmp_path):
+    from muse.examples import copy_model
+
+    copy_model("default_new_input", tmp_path)
+    return tmp_path / "model"
+
+
+@fixture
+def con():
+    return duckdb.connect(":memory:")
+
+
+@fixture
+def populate_regions(default_new_input, con):
+    from muse.new_input.readers import read_regions_csv
+
+    with open(default_new_input / "regions.csv") as f:
+        return read_regions_csv(f, con)
+
+
+@fixture
+def populate_commodities(default_new_input, con):
+    from muse.new_input.readers import read_commodities_csv
+
+    with open(default_new_input / "commodities.csv") as f:
+        return read_commodities_csv(f, con)
+
+
+@fixture
+def populate_demand(default_new_input, con, populate_regions, populate_commodities):
+    from muse.new_input.readers import read_demand_csv
+
+    with open(default_new_input / "demand.csv") as f:
+        return read_demand_csv(f, con)
+
+
+def test_read_regions(populate_regions):
+    assert populate_regions["name"] == np.array(["R1"])
+
+
+def test_read_new_global_commodities(populate_commodities):
+    data = populate_commodities
+    assert list(data["name"]) == ["electricity", "gas", "heat", "wind", "CO2f"]
+    assert list(data["type"]) == ["energy"] * 5
+    assert list(data["unit"]) == ["PJ"] * 4 + ["kt"]
+
+
+def test_calculate_global_commodities(populate_commodities):
+    from muse.new_input.readers import calculate_global_commodities
+
+    data = calculate_global_commodities(populate_commodities)
+
+    assert isinstance(data, xr.Dataset)
+    assert set(data.dims) == {"commodity"}
+    for dt in data.dtypes.values():
+        assert np.issubdtype(dt, np.dtype("str"))
+
+    assert list(data.coords["commodity"].values) == list(populate_commodities["name"])
+    assert list(data.data_vars["type"].values) == list(populate_commodities["type"])
+    assert list(data.data_vars["unit"].values) == list(populate_commodities["unit"])
+
+
+def test_read_new_global_commodities_type_constraint(default_new_input, con):
+    from muse.new_input.readers import read_commodities_csv
+
+    csv = StringIO("name,type,unit\nfoo,invalid,bar\n")
+    with raises(duckdb.ConstraintException):
+        read_commodities_csv(csv, con)
+
+
+def test_new_read_demand_csv(populate_demand):
+    data = populate_demand
+    assert np.all(data["year"] == np.array([2020, 2050]))
+    assert np.all(data["commodity"] == np.array(["heat", "heat"]))
+    assert np.all(data["region"] == np.array(["R1", "R1"]))
+    assert np.all(data["demand"] == np.array([10, 30]))
+
+
+def test_new_read_demand_csv_commodity_constraint(
+    default_new_input, con, populate_commodities, populate_regions
+):
+    from muse.new_input.readers import read_demand_csv
+
+    csv = StringIO("year,commodity_name,region,demand\n2020,invalid,R1,0\n")
+    with raises(duckdb.ConstraintException, match=".*foreign key.*"):
+        read_demand_csv(csv, con)
+
+
+def test_new_read_demand_csv_region_constraint(
+    default_new_input, con, populate_commodities, populate_regions
+):
+    from muse.new_input.readers import read_demand_csv
+
+    csv = StringIO("year,commodity_name,region,demand\n2020,heat,invalid,0\n")
+    with raises(duckdb.ConstraintException, match=".*foreign key.*"):
+        read_demand_csv(csv, con)
+
+
+@mark.xfail
+def test_demand_dataset(default_new_input):
+    import duckdb
+    from muse.new_input.readers import read_commodities, read_demand, read_regions
+
+    con = duckdb.connect(":memory:")
+
+    read_regions(default_new_input, con)
+    read_commodities(default_new_input, con)
+    data = read_demand(default_new_input, con)
+
+    assert isinstance(data, xr.DataArray)
+    assert data.dtype == np.float64
+
+    assert set(data.dims) == {"year", "commodity", "region", "timeslice"}
+    assert list(data.coords["region"].values) == ["R1"]
+    assert list(data.coords["timeslice"].values) == list(range(1, 7))
+    assert list(data.coords["year"].values) == [2020, 2050]
+    assert set(data.coords["commodity"].values) == {
+        "electricity",
+        "gas",
+        "heat",
+        "wind",
+        "CO2f",
+    }
+
+    assert data.sel(year=2020, commodity="electricity", region="R1", timeslice=0) == 1
+
+
+@mark.xfail
+def test_new_read_initial_market(default_new_input):
+    from muse.new_input.readers import read_inputs
+
+    all_data = read_inputs(default_new_input)
+    data = all_data["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,