First pass at duckdb data interface

Author: cc-a

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

@@ -1,6 +1,6 @@
-commodity_name,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,9 @@
+from io import StringIO
 from itertools import chain, permutations
 from pathlib import Path
 from unittest.mock import patch
 
+import duckdb
 import numpy as np
 import toml
 import xarray as xr
@@ -861,40 +863,110 @@ def default_new_input(tmp_path):
     from muse.examples import copy_model
 
     copy_model("default_new_input", tmp_path)
-    return tmp_path
+    return tmp_path / "model"
 
 
-@mark.xfail
-def test_read_new_global_commodities(default_new_input):
-    from muse.new_input.readers import read_inputs
+@fixture
+def con():
+    return duckdb.connect(":memory:")
 
-    all_data = read_inputs(default_new_input)
-    data = all_data["global_commodities"]
+
+@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"}
-    assert dict(data.dtypes) == dict(
-        type=np.dtype("str"),
-        unit=np.dtype("str"),
-    )
+    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)
 
-    assert list(data.coords["commodity"].values) == [
-        "electricity",
-        "gas",
-        "heat",
-        "wind",
-        "CO2f",
-    ]
-    assert list(data.data_vars["type"].values) == ["energy"] * 5
-    assert list(data.data_vars["unit"].values) == ["PJ"] * 4 + ["kt"]
+
+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_read_demand(default_new_input):
-    from muse.new_input.readers import read_inputs
+def test_demand_dataset(default_new_input):
+    import duckdb
+    from muse.new_input.readers import read_commodities, read_demand, read_regions
 
-    all_data = read_inputs(default_new_input)
-    data = all_data["demand"]
+    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