Update tables for new csv columns

Author: tsmbland

src/muse/new_input/readers.py

@@ -31,21 +31,21 @@ def read_inputs(data_dir):
 
 def read_commodities_csv(buffer_, con):
     sql = """CREATE TABLE commodities (
-      name VARCHAR PRIMARY KEY,
+      id 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()
+    con.sql("INSERT INTO commodities SELECT id, type, unit FROM rel;")
+    return con.sql("select * from commodities").fetchnumpy()
 
 
 def read_commodity_trade_csv(buffer_, con):
     sql = """CREATE TABLE commodity_trade (
-    commodity VARCHAR REFERENCES commodities(name),
-    region VARCHAR REFERENCES regions(name),
+    commodity VARCHAR REFERENCES commodities(id),
+    region VARCHAR REFERENCES regions(id),
     year BIGINT,
     import DOUBLE,
     export DOUBLE,
@@ -54,68 +54,68 @@ def read_commodity_trade_csv(buffer_, con):
     con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
     con.sql("""INSERT INTO commodity_trade SELECT
-            commodity, region, year, import, export FROM rel;""")
+            commodity_id, region_id, year, import, export FROM rel;""")
     return con.sql("SELECT * from commodity_trade").fetchnumpy()
 
 
 def read_commodity_costs_csv(buffer_, con):
     sql = """CREATE TABLE commodity_costs (
+    commodity VARCHAR REFERENCES commodities(id),
+    region VARCHAR REFERENCES regions(id),
     year BIGINT,
-    region VARCHAR REFERENCES regions(name),
-    commodity VARCHAR REFERENCES commodities(name),
     value DOUBLE,
     );
     """
     con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
     con.sql("""INSERT INTO commodity_costs SELECT
-            year, region, commodity_name, value FROM rel;""")
+            commidity_id, region_id, year, value FROM rel;""")
     return con.sql("SELECT * from commodity_costs").fetchnumpy()
 
 
 def read_demand_csv(buffer_, con):
     sql = """CREATE TABLE demand (
+    commodity VARCHAR REFERENCES commodities(id),
+    region VARCHAR REFERENCES regions(id),
     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;")
+    con.sql("INSERT INTO demand SELECT commodity_id, region_id, year, demand FROM rel;")
     return con.sql("SELECT * from demand").fetchnumpy()
 
 
 def read_demand_slicing_csv(buffer_, con):
     sql = """CREATE TABLE demand_slicing (
-    commodity VARCHAR REFERENCES commodities(name),
-    region VARCHAR REFERENCES regions(name),
+    commodity VARCHAR REFERENCES commodities(id),
+    region VARCHAR REFERENCES regions(id),
+    year BIGINT,
     timeslice VARCHAR,
     fraction DOUBLE CHECK (fraction >= 0 AND fraction <= 1),
-    year BIGINT,
     );
     """
     con.sql(sql)
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
     con.sql("""INSERT INTO demand_slicing SELECT
-            commodity, region, timeslice, fraction, year FROM rel;""")
+            commodity_id, region_id, year, timeslice, fraction FROM rel;""")
     return con.sql("SELECT * from demand_slicing").fetchnumpy()
 
 
 def read_regions_csv(buffer_, con):
     sql = """CREATE TABLE regions (
-      name VARCHAR PRIMARY KEY,
+      id 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()
+    con.sql("INSERT INTO regions SELECT id FROM rel;")
+    return con.sql("SELECT * from regions").fetchnumpy()
 
 
 def calculate_global_commodities(commodities):
-    names = commodities["name"].astype(np.dtype("str"))
+    names = commodities["id"].astype(np.dtype("str"))
     types = commodities["type"].astype(np.dtype("str"))
     units = commodities["unit"].astype(np.dtype("str"))
 

tests/test_readers.py

@@ -332,14 +332,6 @@ 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
@@ -356,13 +348,21 @@ def populate_demand(default_new_input, con, populate_regions, populate_commoditi
         return read_demand_csv(f, con)
 
 
+@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)
+
+
 def test_read_regions(populate_regions):
-    assert populate_regions["name"] == np.array(["R1"])
+    assert populate_regions["id"] == 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["id"]) == ["electricity", "gas", "heat", "wind", "CO2f"]
     assert list(data["type"]) == ["energy"] * 5
     assert list(data["unit"]) == ["PJ"] * 4 + ["kt"]
 
@@ -377,15 +377,15 @@ def test_calculate_global_commodities(populate_commodities):
     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.coords["commodity"].values) == list(populate_commodities["id"])
     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")
+    csv = StringIO("id,type,unit\nfoo,invalid,bar\n")
     with raises(duckdb.ConstraintException):
         read_commodities_csv(csv, con)
 
@@ -403,7 +403,7 @@ def test_new_read_demand_csv_commodity_constraint(
 ):
     from muse.new_input.readers import read_demand_csv
 
-    csv = StringIO("year,commodity_name,region,demand\n2020,invalid,R1,0\n")
+    csv = StringIO("year,commodity_id,region_id,demand\n2020,invalid,R1,0\n")
     with raises(duckdb.ConstraintException, match=".*foreign key.*"):
         read_demand_csv(csv, con)
 
@@ -413,7 +413,7 @@ def test_new_read_demand_csv_region_constraint(
 ):
     from muse.new_input.readers import read_demand_csv
 
-    csv = StringIO("year,commodity_name,region,demand\n2020,heat,invalid,0\n")
+    csv = StringIO("year,commodity_id,region_id,demand\n2020,heat,invalid,0\n")
     with raises(duckdb.ConstraintException, match=".*foreign key.*"):
         read_demand_csv(csv, con)