From 5e4b5df9235a7767f1b7db1805d2de18a5057dc3 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 09:06:13 +0100 Subject: [PATCH 01/33] Tidy test_agents --- tests/test_agents.py | 156 ++++++++++++++++--------------------------- 1 file changed, 57 insertions(+), 99 deletions(-) diff --git a/tests/test_agents.py b/tests/test_agents.py index d1b97b44..9507c9a8 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -15,134 +15,98 @@ def assets(): result = Dataset() result["year"] = "year", range(2010, 2031) assets = set((randint(2010, 2020), choice(list("technology"))) for i in range(30)) - assets = list(assets) result["installed"] = "asset", [u[0] for u in assets] result["technology"] = "asset", [u[1] for u in assets] shape = len(result.year), len(result.asset) result["capacity"] = ("year", "asset"), ones(shape, dtype="int") - result["maxyears"] = "asset", randint(2010, 2030, len(result.asset)) - result["capa"] = "asset", randint(1, 100, len(result.asset)) - result["capacity"] = result.capacity.where(result.year <= result.maxyears, 0) - result["capacity"] *= result.capa - result = result.drop_vars(("maxyears", "capa")) + result["capacity"] *= randint(1, 100, len(result.asset)) + + # Generate max years with correct shape for broadcasting + max_years = randint(2010, 2030, len(result.asset)) + result["capacity"] = result.capacity.where( + result.year.values[:, None] <= max_years, 0 + ) result = result.set_coords(("installed", "technology")) return result.sel(year=avoid_repetitions(result.capacity)) -def test_create_retrofit(agent_args, technologies, stock): +@fixture +def create_test_agent(): + """Helper fixture to create and test agents with common assertions.""" from muse.agents.agent import Agent from muse.agents.factories import create_agent + def _create_and_test( + agent_type, technologies, stock, agent_args, year=2010, **kwargs + ): + agent = create_agent( + agent_type=agent_type, + technologies=technologies, + capacity=stock.capacity, + year=year, + **{**agent_args, **kwargs}, + ) + assert isinstance(agent, Agent) + assert "asset" in agent.assets.dims + assert "year" in agent.assets.dims + assert "region" not in agent.assets.dims + assert "commodity" not in agent.assets.dims + return agent + + return _create_and_test + + +def test_create_retrofit(agent_args, technologies, stock, create_test_agent): + # Test with zero share agent_args["share"] = "agent_share_zero" - agent = create_agent( - agent_type="Retrofit", - technologies=technologies, - capacity=stock.capacity, - year=2010, - **agent_args, - ) - assert isinstance(agent, Agent) + agent = create_test_agent("Retrofit", technologies, stock, agent_args) assert (agent.assets.capacity == 0).all() - assert "asset" in agent.assets.dims and len(agent.assets.asset) != 0 - assert "year" in agent.assets.dims or len(agent.assets.year) > 1 - assert "region" not in agent.assets.dims - assert "commodity" not in agent.assets.dims + assert len(agent.assets.asset) != 0 + # Test with non-zero share agent_args["share"] = "agent_share" - agent = create_agent( - agent_type="Retrofit", - technologies=technologies, - capacity=stock.capacity, - year=2010, - **agent_args, - ) - assert isinstance(agent, Agent) - assert "asset" in agent.assets.dims - assert len(agent.assets.capacity) != 0 + agent = create_test_agent("Retrofit", technologies, stock, agent_args) assert (agent.assets.capacity != 0).any() -def test_create_newcapa(agent_args, technologies, stock): - from muse.agents.agent import Agent - from muse.agents.factories import create_agent - - # If there is no retrofit, new capa should behave identical to retrofit. - agent_args["share"] = "agent_share_zero" - agent_args["retrofit_present"] = False - agent = create_agent( - agent_type="Newcapa", - technologies=technologies, - capacity=stock.capacity, - year=2010, - **agent_args, - ) - assert isinstance(agent, Agent) +def test_create_newcapa(agent_args, technologies, stock, create_test_agent): + # Test without retrofit + agent_args.update({"share": "agent_share_zero", "retrofit_present": False}) + agent = create_test_agent("Newcapa", technologies, stock, agent_args) assert (agent.assets.capacity == 0).all() - assert "asset" in agent.assets.dims and len(agent.assets.asset) != 0 - assert "year" in agent.assets.dims or len(agent.assets.year) > 1 - assert "region" not in agent.assets.dims - assert "commodity" not in agent.assets.dims assert agent.merge_transform.__name__ == "merge" + # Test with non-zero share agent_args["share"] = "agent_share" - agent = create_agent( - agent_type="Newcapa", - technologies=technologies, - capacity=stock.capacity, - year=2010, - **agent_args, - ) - assert isinstance(agent, Agent) - assert "asset" in agent.assets.dims - assert len(agent.assets.capacity) != 0 + agent = create_test_agent("Newcapa", technologies, stock, agent_args) assert (agent.assets.capacity != 0).any() assert agent.merge_transform.__name__ == "merge" - # If there are retrofit agents, these are really newcapa agents with no capacity - agent_args["share"] = "agent_share" - agent_args["retrofit_present"] = True - agent = create_agent( - agent_type="Newcapa", - technologies=technologies, - capacity=stock.capacity, - year=2010, - **agent_args, - ) - - assert isinstance(agent, Agent) - assert "asset" in agent.assets.dims - assert len(agent.assets.capacity) != 0 + # Test with retrofit present + agent_args.update({"share": "agent_share", "retrofit_present": True}) + agent = create_test_agent("Newcapa", technologies, stock, agent_args) assert (agent.assets.capacity == 0).all() assert agent.merge_transform.__name__ == "new" -def test_issue_835_and_842(agent_args, technologies, stock): - from muse.agents.agent import Agent - from muse.agents.factories import create_agent - +def test_issue_835_and_842(agent_args, technologies, stock, create_test_agent): agent_args["share"] = "agent_share_zero" - agent = create_agent( - agent_type="Retrofit", - technologies=technologies, - capacity=stock.capacity, - search_rules="from_techs->compress", - year=2010, - **agent_args, + agent = create_test_agent( + "Retrofit", technologies, stock, agent_args, search_rules="from_techs->compress" ) - assert isinstance(agent, Agent) assert (agent.assets.capacity == 0).all() - assert "asset" in agent.assets.dims and len(agent.assets.asset) != 0 - assert "year" in agent.assets.dims or len(agent.assets.year) > 1 - assert "region" not in agent.assets.dims - assert "commodity" not in agent.assets.dims + assert len(agent.assets.asset) != 0 @mark.xfail(reason="Retrofit agents will be deprecated.") def test_run_retro_agent(retro_agent, technologies, agent_market, demand_share): - # make sure capacity limits are not reached - technologies.total_capacity_limit[:] = retro_agent.assets.capacity.sum() * 100 - technologies.max_capacity_addition[:] = retro_agent.assets.capacity.sum() * 100 - technologies.max_capacity_growth[:] = retro_agent.assets.capacity.sum() * 100 + capacity_multiplier = retro_agent.assets.capacity.sum() * 100 + for attr in [ + "total_capacity_limit", + "max_capacity_addition", + "max_capacity_growth", + ]: + setattr(technologies, attr, capacity_multiplier) investment_year = int(agent_market.year[1]) retro_agent.next( @@ -162,15 +126,11 @@ def test_merge_assets(assets): n = len(assets.asset) current = assets.sel(asset=range(n - 2)) current = current.sel(year=avoid_repetitions(current.capacity)) - new = assets.sel(asset=range(n - 2, n)) new = new.sel(year=avoid_repetitions(new.capacity)) actual = merge_assets(current, new) - - multi_assets = coords_to_multiindex(assets) - multi_actual = coords_to_multiindex(actual) - assert (multi_actual == multi_assets).all() + assert (coords_to_multiindex(actual) == coords_to_multiindex(assets)).all() def test_clean_assets(assets): @@ -186,14 +146,12 @@ def test_clean_assets(assets): cleaned = clean_assets(assets, current_year) assert (cleaned.year >= current_year).all() - # fmt: disable empties = set( zip( assets.sel(asset=iempties).technology.values, assets.sel(asset=iempties).installed.values, ) ) - # fmt: enable cleanies = set(zip(cleaned.technology.values, cleaned.installed.values)) originals = set(zip(assets.technology.values, assets.installed.values)) assert empties.isdisjoint(cleanies) From 92d878ef6a499db51eb0e6eb6bd1c8e41e4f8928 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 09:16:36 +0100 Subject: [PATCH 02/33] Tidy test_demand_matching --- tests/test_demand_matching.py | 238 +++++++++++++++++----------------- 1 file changed, 120 insertions(+), 118 deletions(-) diff --git a/tests/test_demand_matching.py b/tests/test_demand_matching.py index 4cf25425..414a6180 100644 --- a/tests/test_demand_matching.py +++ b/tests/test_demand_matching.py @@ -1,12 +1,20 @@ +import xarray as xr +from numpy import arange +from numpy.random import choice, randint, random from pytest import approx, fixture +from muse.demand_matching import demand_matching + + +def assert_matches_demand(result, demand): + """Helper function to check if result matches demand after broadcasting.""" + actual, expected = xr.broadcast(result, demand) + assert actual.data == approx(expected.data) + @fixture(params=["without c", "with c"]) def demand(request): - from numpy.random import choice, randint - from xarray import DataArray - - demand = DataArray( + demand = xr.DataArray( randint(0, 10, (5, 3, 4)), coords={ "a": choice(range(15), 5, replace=False), @@ -15,178 +23,172 @@ def demand(request): }, dims=("a", "b", "c"), ) - if request.param == "without c": - demand = demand.sum("c") - return demand.astype("float64") + return demand.sum("c") if request.param == "without c" else demand.astype("float64") @fixture def cost(demand): - from numpy.random import randint - from xarray import DataArray - - return DataArray( + return xr.DataArray( randint(0, 10, size=demand.shape), coords=demand.coords, dims=demand.dims ) @fixture def dataset(timeslice): - """cost, demand, max_production, demand matches exactly max production.""" - from numpy.random import choice, randint, random - from xarray import Dataset + """Creates a test dataset with cost, demand, and max_production.""" + ds = xr.Dataset() - dataset = Dataset() + # Create dimension coordinates for dim in ("d", "m", "c", "dc", "dm", "cm", "dcm"): nitems = randint(1, 10) - dataset[dim] = dim, choice(list(range(2 * nitems)), nitems, replace=False) + ds[dim] = dim, choice(range(2 * nitems), nitems, replace=False) - shape = (len(dataset.c), len(dataset.dc), len(dataset.cm), len(dataset.dcm)) - dataset["cost"] = (("c", "dc", "cm", "dcm"), randint(0, 5, shape).astype("float64")) + # Set up cost array + shape_cost = (len(ds.c), len(ds.dc), len(ds.cm), len(ds.dcm)) + ds["cost"] = (("c", "dc", "cm", "dcm"), randint(0, 5, shape_cost).astype("float64")) - shape = (len(dataset.m), len(dataset.dm), len(dataset.cm), len(dataset.dcm)) - dataset["max_production"] = ( + # Set up max_production with some zero values + shape_prod = (len(ds.m), len(ds.dm), len(ds.cm), len(ds.dcm)) + ds["max_production"] = ( ("m", "dm", "cm", "dcm"), - 1.0 * randint(0, 10, shape) * (0 == randint(0, 2, shape)), + 1.0 * randint(0, 10, shape_prod) * (0 == randint(0, 2, shape_prod)), ) - summed_production = dataset.max_production.sum(("m", "cm")) - dc = dataset.dc.copy(data=random(dataset.dc.shape)) - dc_share = dc / dc.sum() - d = dataset.d.copy(data=random(dataset.d.shape)) - d_share = d / d.sum() - dataset["demand"] = summed_production * dc_share * d_share - return dataset + # Calculate demand based on production + summed_production = ds.max_production.sum(("m", "cm")) + # Create shares with correct shapes + dc_share = xr.DataArray(random(len(ds.dc)), coords={"dc": ds.dc}, dims="dc") + dc_share = dc_share / dc_share.sum() -def test_cost_order(dataset): - from numpy import arange - from xarray import broadcast + d_share = xr.DataArray(random(len(ds.d)), coords={"d": ds.d}, dims="d") + d_share = d_share / d_share.sum() - from muse.demand_matching import demand_matching + # Broadcast and multiply + ds["demand"] = summed_production * dc_share * d_share - # simplify dataset. bail out if size is too small for test. + return ds + + +def test_cost_order(dataset): + """Tests the ordering of technology selection based on cost. + + This test verifies that: + 1. When costs are ordered, the lowest cost technology is selected first + 2. When two technologies have equal costs, they share the demand equally + 3. When a technology becomes more expensive, it is not selected + """ + # Simplify dataset to focus on technology selection dimensions ds = dataset.sum(set(dataset.dims).difference(("dm", "cm"))) if ds.cm.size < 2: - return + return # Skip test if insufficient technologies for comparison - # ensure we know which tech is selected and that any one tech can fulfill the whole - # demand. - ds.cost[:] = arange(ds.cost.size).reshape(ds.cost.shape) - ds.max_production[:] = ds.demand.sum() + # Set up initial conditions where any tech can fulfill total demand + total_demand = ds.demand.sum() + ds.max_production[:] = total_demand - # should select first tech + # Test Case 1: Sequential cost ordering - should select first tech + ds.cost[:] = arange(ds.cost.size).reshape(ds.cost.shape) result = demand_matching(ds.demand, ds.cost) - actual, dems = broadcast(result.sum(ds.cost.dims), ds.demand) - assert actual.data == approx(dems.data) - assert result.sum(ds.demand.dims)[0].data == approx(ds.demand.sum().data) + assert_matches_demand(result.sum(ds.cost.dims), ds.demand) + assert result.sum(ds.demand.dims)[0].data == approx(total_demand.data) - # should select first and second tech + # Test Case 2: Equal costs for first two techs - should split demand equally ds.cost[0] = 1 result = demand_matching(ds.demand, ds.cost) - actual, dems = broadcast(result.sum(ds.cost.dims), ds.demand) - assert actual.data == approx(dems.data) + assert_matches_demand(result.sum(ds.cost.dims), ds.demand) summed = result.sum(ds.demand.dims).data - assert summed[0] == approx(summed[1]) - assert 2 * summed[0] == approx(ds.demand.sum().data) + assert summed[0] == approx(summed[1]), ( + "First two technologies should share demand equally" + ) + assert 2 * summed[0] == approx(total_demand.data), ( + "Sum of shared demand should equal total" + ) - # should select second tech + # Test Case 3: First tech more expensive - should select second tech only ds.cost[0] = 2 result = demand_matching(ds.demand, ds.cost) - actual, dems = broadcast(result.sum(ds.cost.dims), ds.demand) - assert actual.data == approx(dems.data) + assert_matches_demand(result.sum(ds.cost.dims), ds.demand) summed = result.sum(ds.demand.dims).data - assert summed[1] == approx(ds.demand.sum().data) + assert summed[1] == approx(total_demand.data), ( + "Second technology should fulfill all demand" + ) def test_no_constraints_no_i_dims(demand, cost): - from xarray import broadcast - - from muse.demand_matching import demand_matching + # Test without b dimension in cost + result = demand_matching(demand, cost.sum("b")) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result, demand) - x = demand_matching(demand, cost.sum("b")) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x, demand) - assert x.values == approx(expected.values) + # Test with all dimensions + result = demand_matching(demand, cost) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result, demand) - x = demand_matching(demand, cost) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x, demand) - assert x.values == approx(expected.values) - - x = demand_matching(demand.sum("a"), cost) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x.sum("a"), demand.sum("a")) - assert x.values == approx(expected.values) + # Test with summed a dimension + result = demand_matching(demand.sum("a"), cost) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result.sum("a"), demand.sum("a")) def test_no_constraints_i_dims(demand, cost): - from xarray import broadcast - - from muse.demand_matching import demand_matching - - x = demand_matching(demand.sum("a"), cost, protected_dims={"a"}) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x.sum("a"), demand.sum("a")) - assert x.values == approx(expected.values) + # Test protected dimension 'a' + result = demand_matching(demand.sum("a"), cost, protected_dims={"a"}) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result.sum("a"), demand.sum("a")) - x = demand_matching(demand.sum("b"), cost, protected_dims={"b"}) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x.sum("b"), demand.sum("b")) - assert x.values == approx(expected.values) + # Test protected dimension 'b' + result = demand_matching(demand.sum("b"), cost, protected_dims={"b"}) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result.sum("b"), demand.sum("b")) def test_one_non_binding_constraint(demand, cost): - """Constraint where excess is always 0.""" - from xarray import broadcast + """Tests constraints where excess is always 0.""" + # Test with full dimensions + result = demand_matching(demand, cost, 2 * demand) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result, demand) - from muse.demand_matching import demand_matching - - x = demand_matching(demand, cost, 2 * demand) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x, demand) - assert x.values == approx(expected.values) - - x = demand_matching(demand, cost, (2 * demand).sum("a")) - assert set(x.dims) == set(demand.dims) - x, expected = broadcast(x, demand) - assert x.values == approx(expected.values) + # Test with summed dimension + result = demand_matching(demand, cost, (2 * demand).sum("a")) + assert set(result.dims) == set(demand.dims) + assert_matches_demand(result, demand) def test_one_cutting_constraint(demand, cost): - """Constraint where excess is not always 0.""" - from xarray import broadcast - - from muse.demand_matching import demand_matching - + """Tests constraints where excess is not always 0.""" constraint = demand.sum("a") * 0.5 - x = demand_matching(demand, cost, constraint) - assert set(x.dims) == set(demand.dims) - assert (x.sum("a") - constraint <= 1e-12).all() - expected, actual = broadcast(x.sum("a") + constraint, demand.sum("a")) + # Test with full demand + result = demand_matching(demand, cost, constraint) + assert set(result.dims) == set(demand.dims) + assert (result.sum("a") - constraint <= 1e-12).all() + expected, actual = xr.broadcast(result.sum("a") + constraint, demand.sum("a")) assert actual.values == approx(expected.values) - x = demand_matching(demand.sum("b"), cost, constraint) - assert set(x.dims) == set(demand.dims) - assert (x.sum("b") - demand.sum("b") < 1e-12).all() + # Test with summed demand + result = demand_matching(demand.sum("b"), cost, constraint) + assert set(result.dims) == set(demand.dims) + assert (result.sum("b") - demand.sum("b") < 1e-12).all() def test_two_cutting_constraint(demand, cost): - """Constraint where excess is not always 0.""" - from muse.demand_matching import demand_matching - + """Tests multiple constraints where excess is not always 0.""" constraint0 = demand.sum("a") * 0.75 constraint1 = demand.sum("b") * 0.85 + demand.sum("a") * 0.80 - x = demand_matching(demand, cost, constraint0, constraint1) - assert set(x.dims) == set(demand.dims) - assert (x.sum("a") - constraint0 <= 1e-12).all() - assert (x - constraint1 <= 1e-12).all() - assert (x - demand <= 1e-12).all() - - x = demand_matching(demand.sum("a"), cost, constraint0, constraint1) - assert set(x.dims) == set(demand.dims) - assert (x.sum("a") - constraint0 <= 1e-12).all() - assert (x - constraint1 <= 1e-12).all() - assert (x.sum("a") - demand.sum("a") <= 1e-12).all() + # Test with full demand + result = demand_matching(demand, cost, constraint0, constraint1) + assert set(result.dims) == set(demand.dims) + assert (result.sum("a") - constraint0 <= 1e-12).all() + assert (result - constraint1 <= 1e-12).all() + assert (result - demand <= 1e-12).all() + + # Test with summed demand + result = demand_matching(demand.sum("a"), cost, constraint0, constraint1) + assert set(result.dims) == set(demand.dims) + assert (result.sum("a") - constraint0 <= 1e-12).all() + assert (result - constraint1 <= 1e-12).all() + assert (result.sum("a") - demand.sum("a") <= 1e-12).all() From 885670e80644754a2da14a5dd1628538264996f7 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 09:24:12 +0100 Subject: [PATCH 03/33] Tidy test_costs --- tests/test_costs.py | 165 +++++++++++++------------------------------- 1 file changed, 48 insertions(+), 117 deletions(-) diff --git a/tests/test_costs.py b/tests/test_costs.py index 8d032cff..f66237b0 100644 --- a/tests/test_costs.py +++ b/tests/test_costs.py @@ -1,5 +1,25 @@ from numpy import isclose, isfinite from pytest import fixture, mark, raises +from xarray.testing import assert_allclose + +from muse.costs import ( + annual_to_lifetime, + capital_costs, + capital_recovery_factor, + environmental_costs, + equivalent_annual_cost, + fixed_costs, + fuel_costs, + levelized_cost_of_energy, + material_costs, + net_present_cost, + net_present_value, + running_costs, + supply_cost, + variable_costs, +) +from muse.quantities import production_amplitude +from muse.timeslices import broadcast_timeslice YEAR = 2030 @@ -9,10 +29,7 @@ def _capacity(_technologies, demand_share): """Capacity for each asset.""" from muse.quantities import capacity_to_service_demand - capacity = capacity_to_service_demand( - technologies=_technologies, demand=demand_share - ) - return capacity + return capacity_to_service_demand(technologies=_technologies, demand=demand_share) @fixture @@ -37,12 +54,11 @@ def _production(_technologies, _capacity): """Production data for each asset.""" from muse.timeslices import broadcast_timeslice, distribute_timeslice - production = ( + return ( broadcast_timeslice(_capacity) * distribute_timeslice(_technologies.fixed_outputs) * broadcast_timeslice(_technologies.utilization_factor) ) - return production @fixture @@ -50,75 +66,53 @@ def _consumption(_technologies, _capacity): """Consumption data for each asset.""" from muse.timeslices import broadcast_timeslice, distribute_timeslice - consumption = ( + return ( broadcast_timeslice(_capacity) * distribute_timeslice(_technologies.fixed_inputs) * broadcast_timeslice(_technologies.utilization_factor) ) - return consumption def test_fixtures(_technologies, _prices, _capacity, _production, _consumption): - """Validating that the fixtures have appropriate dimensions.""" + """Validate fixture dimensions.""" assert set(_technologies.dims) == {"asset", "commodity"} assert set(_prices.dims) == {"asset", "commodity", "timeslice"} assert set(_capacity.dims) == {"asset"} - assert ( - set(_production.dims) - == set(_consumption.dims) - == { - "asset", - "commodity", - "timeslice", - } - ) + assert set(_production.dims) == {"asset", "commodity", "timeslice"} + assert set(_consumption.dims) == {"asset", "commodity", "timeslice"} def test_capital_costs(_technologies, _capacity): - from muse.costs import capital_costs - result = capital_costs(_technologies, _capacity) assert set(result.dims) == {"asset"} def test_environmental_costs(_technologies, _prices, _production): - from muse.costs import environmental_costs - result = environmental_costs(_technologies, _prices, _production) assert set(result.dims) == {"asset", "timeslice"} def test_fuel_costs(_technologies, _prices, _consumption): - from muse.costs import fuel_costs - result = fuel_costs(_technologies, _prices, _consumption) assert set(result.dims) == {"asset", "timeslice"} def test_material_costs(_technologies, _prices, _consumption): - from muse.costs import material_costs - result = material_costs(_technologies, _prices, _consumption) assert set(result.dims) == {"asset", "timeslice"} def test_fixed_costs(_technologies, _capacity): - from muse.costs import fixed_costs - result = fixed_costs(_technologies, _capacity) assert set(result.dims) == {"asset"} def test_variable_costs(_technologies, _production): - from muse.costs import variable_costs - result = variable_costs(_technologies, _production) assert set(result.dims) == {"asset"} def test_running_costs(_technologies, _prices, _capacity, _production, _consumption): - from muse.costs import running_costs - result = running_costs(_technologies, _prices, _capacity, _production, _consumption) assert set(result.dims) == {"asset", "timeslice"} @@ -126,8 +120,6 @@ def test_running_costs(_technologies, _prices, _capacity, _production, _consumpt def test_net_present_value( _technologies, _prices, _capacity, _production, _consumption ): - from muse.costs import net_present_value - result = net_present_value( _technologies, _prices, _capacity, _production, _consumption ) @@ -135,8 +127,6 @@ def test_net_present_value( def test_net_present_cost(_technologies, _prices, _capacity, _production, _consumption): - from muse.costs import net_present_cost - result = net_present_cost( _technologies, _prices, _capacity, _production, _consumption ) @@ -146,8 +136,6 @@ def test_net_present_cost(_technologies, _prices, _capacity, _production, _consu def test_equivalent_annual_cost( _technologies, _prices, _capacity, _production, _consumption ): - from muse.costs import equivalent_annual_cost - result = equivalent_annual_cost( _technologies, _prices, _capacity, _production, _consumption ) @@ -158,8 +146,6 @@ def test_equivalent_annual_cost( def test_levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method ): - from muse.costs import levelized_cost_of_energy - result = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method=method ) @@ -167,35 +153,24 @@ def test_levelized_cost_of_energy( def test_supply_cost(_technologies, _prices, _capacity, _production, _consumption): - from muse.costs import levelized_cost_of_energy, supply_cost - lcoe = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method="annual" ) result = supply_cost(_production, lcoe) - assert set(result.dims) == { - "commodity", - "region", - "timeslice", - } + assert set(result.dims) == {"commodity", "region", "timeslice"} def test_capital_recovery_factor(_technologies): - from muse.costs import capital_recovery_factor - result = capital_recovery_factor(_technologies) assert set(result.dims) == set(_technologies.interest_rate.dims) - # {"region", "technology"} - # Make sure zero interest rates are supported + # Test zero interest rates _technologies["interest_rate"] = 0 result = capital_recovery_factor(_technologies) assert isfinite(result).all() def test_annual_to_lifetime(_technologies, _prices, _consumption): - from muse.costs import annual_to_lifetime, fuel_costs - _fuel_costs = fuel_costs(_technologies, _prices, _consumption) _fuel_costs_lifetime = annual_to_lifetime(_fuel_costs, _technologies) assert set(_fuel_costs.dims) == set(_fuel_costs_lifetime.dims) @@ -206,31 +181,21 @@ def test_annual_to_lifetime(_technologies, _prices, _consumption): def test_lcoe_flow_scaling( _technologies, _prices, _capacity, _production, _consumption, method ): - """Testing that LCOE is independent of input/output flow scaling. - - In other words, if we change technology flows by a constant factor, the LCOE (which - is a cost per unit of production) should remain unchanged. - - This is a bit more complicated if the variable costs are nonlinear, so we'll set - the exponent to 1 for simplicity. - """ - from muse.costs import levelized_cost_of_energy - + """Test LCOE independence of input/output flow scaling.""" _technologies["var_exp"] = 1 - # LCOE with original inputs + # Original LCOE lcoe1 = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method=method ) - # Scale inputs and outputs by a constant factor -> LCOE should be unchanged - # var_par also needs to be scaled as this relates to units of technology - # activity, not units of commodity consumption/production + # Scale inputs/outputs and var_par by 2 _technologies_scaled = _technologies.copy() - _technologies_scaled["fixed_inputs"] = _technologies["fixed_inputs"] * 2 - _technologies_scaled["flexible_inpits"] = _technologies["flexible_inputs"] * 2 - _technologies_scaled["fixed_outputs"] = _technologies["fixed_outputs"] * 2 - _technologies_scaled["var_par"] = _technologies["var_par"] * 2 + _technologies_scaled["fixed_inputs"] *= 2 + _technologies_scaled["flexible_inputs"] *= 2 + _technologies_scaled["fixed_outputs"] *= 2 + _technologies_scaled["var_par"] *= 2 + lcoe2 = levelized_cost_of_energy( _technologies_scaled, _prices, @@ -246,24 +211,14 @@ def test_lcoe_flow_scaling( def test_lcoe_prod_scaling( _technologies, _prices, _capacity, _production, _consumption, method ): - """Testing that LCOE is independent of production scaling. - - If all costs are linear (exponents = 1), then the LCOE should be independent of - production as long as production, consumption, and capacity are scaled together. - """ - from muse.costs import levelized_cost_of_energy - + """Test LCOE independence of production scaling with linear costs.""" _technologies["var_exp"] = 1 _technologies["cap_exp"] = 1 _technologies["fix_exp"] = 1 - # LCOE with original inputs lcoe1 = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method=method ) - - # Scale consumption, production, and capacity by a constant factor -> LCOE - # should be unchanged lcoe2 = levelized_cost_of_energy( _technologies, _prices, @@ -279,20 +234,14 @@ def test_lcoe_prod_scaling( def test_lcoe_equal_prices( _technologies, _prices, _capacity, _production, _consumption, method ): - """If commodity prices are equal in every timeslice, LCOE should always be equal.""" - from xarray.testing import assert_allclose - - from muse.costs import levelized_cost_of_energy - from muse.timeslices import broadcast_timeslice - - # LCOE with original inputs -> should vary between timeslices + """Test LCOE behavior with uniform prices across timeslices.""" lcoe1 = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method=method ) with raises(AssertionError): assert_allclose(lcoe1, broadcast_timeslice(lcoe1.isel(timeslice=0))) - # LCOE with uniform prices -> should be the same for all timeslices + # Test with uniform prices _prices = broadcast_timeslice(_prices.mean("timeslice")) lcoe2 = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method=method @@ -301,27 +250,17 @@ def test_lcoe_equal_prices( def test_npv_equal_prices(_technologies, _prices, _capacity, _production, _consumption): - """Test NPV with equal commodity prices in every timeslice. - - If commodity prices are equal in every timeslice, NPV should be proportional to - production. - """ - from xarray.testing import assert_allclose - - from muse.costs import net_present_value - from muse.quantities import production_amplitude - from muse.timeslices import broadcast_timeslice - - # NPV with original inputs -> should not be linear with production + """Test NPV linearity with production under uniform prices.""" npv1 = net_present_value( _technologies, _prices, _capacity, _production, _consumption ) tech_activity = production_amplitude(_production, _technologies) npv1_scaled = npv1 / tech_activity + with raises(AssertionError): assert_allclose(npv1_scaled, broadcast_timeslice(npv1_scaled.isel(timeslice=0))) - # NPV with uniform prices -> should be linear with production + # Test with uniform prices _prices = broadcast_timeslice(_prices.mean("timeslice")) npv2 = net_present_value( _technologies, _prices, _capacity, _production, _consumption @@ -334,19 +273,13 @@ def test_npv_equal_prices(_technologies, _prices, _capacity, _production, _consu def test_lcoe_zero_production( _technologies, _prices, _capacity, _production, _consumption, method ): - """If production and consumption are zero, LCOE should always be zero. - - Note: if production/consumption are zero in every timeslice, LCOE is undefined (nan) - """ - from muse.costs import levelized_cost_of_energy - - # LCOE with original inputs + """Test LCOE behavior with zero production.""" lcoe1 = levelized_cost_of_energy( _technologies, _prices, _capacity, _production, _consumption, method=method ) assert not (lcoe1.isel(timeslice=0) == 0).all() - # LCOE with zero production/consumption in first timeslice -> LCOE should be zero + # Test with zero production in first timeslice _production.isel(timeslice=0)[:] = 0 _consumption.isel(timeslice=0)[:] = 0 lcoe2 = levelized_cost_of_energy( @@ -359,8 +292,7 @@ def test_lcoe_zero_production( def test_lcoe_aggregate( _technologies, _prices, _capacity, _production, _consumption, method ): - from muse.costs import levelized_cost_of_energy - + """Test LCOE aggregation over timeslices.""" result = levelized_cost_of_energy( _technologies, _prices, @@ -370,12 +302,11 @@ def test_lcoe_aggregate( method=method, aggregate_timeslices=True, ) - assert set(result.dims) == {"asset"} # no timeslice dim + assert set(result.dims) == {"asset"} def test_npv_aggregate(_technologies, _prices, _capacity, _production, _consumption): - from muse.costs import net_present_value - + """Test NPV aggregation over timeslices.""" result = net_present_value( _technologies, _prices, @@ -384,4 +315,4 @@ def test_npv_aggregate(_technologies, _prices, _capacity, _production, _consumpt _consumption, aggregate_timeslices=True, ) - assert set(result.dims) == {"asset"} # no timeslice dim + assert set(result.dims) == {"asset"} From 375b508c2415aa480f4650b55293a9dc5c782680 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 09:31:37 +0100 Subject: [PATCH 04/33] Tidy test_wizard --- tests/test_wizard.py | 179 ++++++++++++++++++------------------------- 1 file changed, 76 insertions(+), 103 deletions(-) diff --git a/tests/test_wizard.py b/tests/test_wizard.py index 0f35125b..2f644945 100644 --- a/tests/test_wizard.py +++ b/tests/test_wizard.py @@ -1,3 +1,5 @@ +from pathlib import Path + import pandas as pd import pytest from tomlkit import dumps, parse @@ -29,172 +31,143 @@ def model_path_retro(tmp_path): return tmp_path / "model" +def assert_values_in_csv(file_path: Path, column: str, expected_values: list): + """Helper function to check if values exist in a CSV column.""" + df = pd.read_csv(file_path) + for value in expected_values: + assert value in df[column].values + + +def assert_columns_exist(file_path: Path, columns: list): + """Helper function to check if columns exist in a CSV file.""" + df = pd.read_csv(file_path) + for column in columns: + assert column in df.columns + + def test_modify_toml(tmp_path): """Test the modify_toml function.""" - # Create a temporary toml file toml_path = tmp_path / "temp.toml" - - # Create initial toml data initial_data = {"name": "Tomm", "age": 299} - - # Write initial toml data to the temporary file toml_path.write_text(dumps(initial_data)) - # Define the function to modify the toml data def modify_function(data): - data["name"] = "Tom" - data["age"] = 29 + data.update({"name": "Tom", "age": 29}) - # Call the modify_toml function modify_toml(toml_path, modify_function) - - # Read the modified toml data modified_data = parse(toml_path.read_text()) - # Assert that the modifications were applied correctly - assert modified_data["name"] == "Tom" - assert modified_data["age"] == 29 + assert modified_data == {"name": "Tom", "age": 29} def test_get_sectors(tmp_path): """Test the get_sectors function.""" - # Create a temporary model folder model_path = tmp_path / "model" model_path.mkdir() - # Create some sector folders with Technodata.csv files - sector1 = model_path / "sector1" - sector1.mkdir(parents=True) - (sector1 / "Technodata.csv").touch() - - sector2 = model_path / "sector2" - sector2.mkdir(parents=True) - (sector2 / "Technodata.csv").touch() - - sector3 = model_path / "sector3" - sector3.mkdir(parents=True) - - # Call the get_sectors function - sectors = get_sectors(model_path) + # Create test sector folders + for sector in ["sector1", "sector2", "sector3"]: + sector_path = model_path / sector + sector_path.mkdir(parents=True) + if sector != "sector3": + (sector_path / "Technodata.csv").touch() - # Check the returned sectors - assert set(sectors) == {"sector1", "sector2"} + assert set(get_sectors(model_path)) == {"sector1", "sector2"} def test_add_new_commodity(model_path): """Test the add_new_commodity function on the default model.""" add_new_commodity(model_path, "new_commodity", "power", "wind") - # Check if the new commodity is added to the global commodities file - global_commodities_file = model_path / "GlobalCommodities.csv" - df = pd.read_csv(global_commodities_file) - assert "new_commodity" in df["CommodityName"].values + # Check global commodities + assert_values_in_csv( + model_path / "GlobalCommodities.csv", "CommodityName", ["new_commodity"] + ) - # Check if the new column is added to additional files + # Check commodity appears in relevant files files_to_check = [ - model_path / file - for file in [ - "power/CommIn.csv", - "power/CommOut.csv", - "Projections.csv", - ] - ] + list((model_path / "residential_presets").glob("*")) + model_path / "power/CommIn.csv", + model_path / "power/CommOut.csv", + model_path / "Projections.csv", + *(model_path / "residential_presets").glob("*"), + ] + for file in files_to_check: - df = pd.read_csv(model_path / file) - assert "new_commodity" in df.columns + assert_columns_exist(file, ["new_commodity"]) def test_add_new_process(model_path): """Test the add_new_process function on the default model.""" add_new_process(model_path, "new_process", "power", "windturbine") - # Check if the new process is added to the files files_to_check = [ - "power/CommIn.csv", - "power/CommOut.csv", - "power/ExistingCapacity.csv", - "power/Technodata.csv", + "CommIn.csv", + "CommOut.csv", + "ExistingCapacity.csv", + "Technodata.csv", ] for file in files_to_check: - df = pd.read_csv(model_path / file) - assert "new_process" in df["ProcessName"].values + assert_values_in_csv( + model_path / "power" / file, "ProcessName", ["new_process"] + ) def test_add_price_data_for_new_year(model_path): """Test the add_price_data_for_new_year function on the default model.""" add_price_data_for_new_year(model_path, "2030", "power", "2020") - # Check if the new price data is added to the files - files_to_check = [ - "power/Technodata.csv", - "power/CommIn.csv", - "power/CommOut.csv", - ] + files_to_check = ["Technodata.csv", "CommIn.csv", "CommOut.csv"] for file in files_to_check: - df = pd.read_csv(model_path / file) - assert "2030" in df["Time"].values + assert_values_in_csv(model_path / "power" / file, "Time", ["2030"]) def test_add_agent(model_path_retro): """Test the add_agent function on the default_retro model.""" add_agent(model_path_retro, "A2", "A1", "Agent3", "Agent4") - # Check if the new agent is added to the Agents.csv file - df = pd.read_csv(model_path_retro / "Agents.csv") - assert "A2" in df["Name"].values - assert "Agent3" in df["AgentShare"].values - assert "Agent4" in df["AgentShare"].values + # Check Agents.csv + assert_values_in_csv(model_path_retro / "Agents.csv", "Name", ["A2"]) + for share in ["Agent3", "Agent4"]: + assert_values_in_csv(model_path_retro / "Agents.csv", "AgentShare", [share]) - # Check if the retrofit agent is added to the Technodata.csv files - sector1_file = model_path_retro / "power/Technodata.csv" - sector2_file = model_path_retro / "gas/Technodata.csv" - df_sector1 = pd.read_csv(sector1_file) - df_sector2 = pd.read_csv(sector2_file) - assert "Agent4" in df_sector1.columns - assert "Agent4" in df_sector2.columns + # Check Technodata.csv files + for sector in ["power", "gas"]: + assert_columns_exist(model_path_retro / sector / "Technodata.csv", ["Agent4"]) def test_add_region(model_path): """Test the add_region function on the default model.""" add_region(model_path, "R2", "R1") - # Check if the new region is added to the settings.toml file + # Check settings.toml with open(model_path / "settings.toml") as f: - modified_settings_data = parse(f.read()) - assert "R2" in modified_settings_data["regions"] - - # Check if the new region is added to the technodata files - sector_files = [ - model_path / sector / file - for sector in get_sectors(model_path) - for file in [ - "Technodata.csv", - "CommIn.csv", - "CommOut.csv", - "ExistingCapacity.csv", - ] + settings = parse(f.read()) + assert "R2" in settings["regions"] + + # Check sector files + files_to_check = [ + "Technodata.csv", + "CommIn.csv", + "CommOut.csv", + "ExistingCapacity.csv", ] - for file in sector_files: - df = pd.read_csv(file) - assert "R2" in df["RegionName"].values + for sector in get_sectors(model_path): + for file in files_to_check: + assert_values_in_csv(model_path / sector / file, "RegionName", ["R2"]) def test_add_timeslice(model_path): """Test the add_timeslice function on the default model.""" add_timeslice(model_path, "midnight", "evening") - # Check if the new timeslice is added to the settings.toml file + # Check settings.toml with open(model_path / "settings.toml") as f: - modified_settings_data = parse(f.read()) - assert "midnight" in modified_settings_data["timeslices"]["all-year"]["all-week"] - n_timeslices = len(modified_settings_data["timeslices"]["all-year"]["all-week"]) - - # Check if the new timeslice is added to the preset files - df_preset1 = pd.read_csv( - model_path / "residential_presets/Residential2020Consumption.csv" - ) - df_preset2 = pd.read_csv( - model_path / "residential_presets/Residential2050Consumption.csv" - ) - assert len(df_preset1["Timeslice"].unique()) == n_timeslices - assert len(df_preset2["Timeslice"].unique()) == n_timeslices + settings = parse(f.read()) + timeslices = settings["timeslices"]["all-year"]["all-week"] + assert "midnight" in timeslices + n_timeslices = len(timeslices) + + # Check preset files + for preset in ["Residential2020Consumption.csv", "Residential2050Consumption.csv"]: + df = pd.read_csv(model_path / "residential_presets" / preset) + assert len(df["Timeslice"].unique()) == n_timeslices From d324bc3ed1fd7d59722d050814cc88493d8b4103 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:14:39 +0100 Subject: [PATCH 05/33] Tidy test_outputs --- tests/test_outputs.py | 472 ++++++++++++++++++++---------------------- 1 file changed, 229 insertions(+), 243 deletions(-) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 10c0ab9d..89b2eb2b 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -13,6 +13,7 @@ @fixture def streetcred(save_registries): + """Create a test output quantity that returns random data.""" from muse.outputs.sector import register_output_quantity @register_output_quantity @@ -28,83 +29,103 @@ def streetcred(*args, **kwargs): ) -@mark.usefixtures("streetcred") -def test_save_with_dir(tmpdir): - from pandas import read_csv +@fixture +def market(): + """Common market fixture used in multiple tests.""" + return xr.DataArray([1], coords={"year": [2010]}, dims="year") + +@fixture +def base_config(tmpdir): + """Common config fixture used in multiple tests.""" path = Path(tmpdir) / "results" / "stuff" - config = { + return { "filename": path / "{Sector}{year}{Quantity}.csv", "quantity": "streetcred", } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here - result = factory(config, sector_name="Yoyo")(market, None, None) - assert len(result) == 1 - assert result[0] == path / "Yoyo2010Streetcred.csv" - assert result[0].exists() - assert result[0].is_file() - read_csv(result[0]) + + +def create_test_data_array(values, coords=None, name="test"): + """Helper function to create test DataArrays.""" + if coords is None: + coords = dict(a=[2, 4]) + return xr.DataArray(values, coords=coords, dims="a", name=name) + + +def assert_file_exists_and_readable(path, expected_columns=None): + """Helper to verify file exists and can be read.""" + assert path.exists() and path.is_file() + df = pd.read_csv(path) + if expected_columns: + assert set(df.columns) == set(expected_columns) + return df @mark.usefixtures("streetcred") -def test_overwrite(tmpdir): - from pytest import raises +def test_save_with_dir(tmpdir, market, base_config): + """Test saving output to directory with sector and year in filename.""" + result = factory(base_config, sector_name="Yoyo")(market, None, None) + assert len(result) == 1 + expected_path = Path(base_config["filename"]).parent / "Yoyo2010Streetcred.csv" + assert result[0] == expected_path + assert_file_exists_and_readable(result[0]) - path = Path(tmpdir) / "results" / "stuff" - config = { - "filename": path / "{Sector}{year}{Quantity}.csv", - "quantity": "streetcred", - } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here - outputter = factory(config, sector_name="Yoyo") + +@mark.usefixtures("streetcred") +def test_overwrite(tmpdir, market, base_config): + """Test file overwrite behavior.""" + outputter = factory(base_config, sector_name="Yoyo") result = outputter(market, None, None) - assert result[0] == path / "Yoyo2010Streetcred.csv" - assert result[0].is_file() + expected_path = Path(base_config["filename"]).parent / "Yoyo2010Streetcred.csv" + assert result[0] == expected_path + assert_file_exists_and_readable(result[0]) # default is to never overwrite with raises(IOError): outputter(market, None, None) - config["overwrite"] = True - factory(config, sector_name="Yoyo")(market, None, None) + base_config["overwrite"] = True + factory(base_config, sector_name="Yoyo")(market, None, None) @mark.usefixtures("streetcred") -def test_save_with_path_to_nc_with_suffix(tmpdir): +@mark.parametrize( + "config_type,suffix", + [ + ("suffix", "nc"), + ("sink", "nc"), + ], +) +def test_save_with_path_to_nc(tmpdir, market, base_config, config_type, suffix): + """Test saving output to NC file with different config types.""" path = Path(tmpdir) / "results" / "stuff" - config = { - "filename": path / "{Sector}{year}{Quantity}{suffix}", - "quantity": "streetcred", - "suffix": "nc", - } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here + if config_type == "suffix": + config = { + "filename": path / "{Sector}{year}{Quantity}{suffix}", + "quantity": "streetcred", + "suffix": suffix, + } + else: + config = { + "filename": path / "{sector}{year}{quantity}.csv", + "quantity": "streetcred", + "sink": suffix, + } result = factory(config, sector_name="Yoyo")(market, None, None) - assert result[0] == path / "Yoyo2010Streetcred.nc" - assert result[0].is_file() - xr.open_dataset(result[0]) - - -@mark.usefixtures("streetcred") -def test_save_with_path_to_nc_with_sink(tmpdir): - path = Path(tmpdir) / "results" / "stuff" - # can use None because we **know** none of the arguments are used here - config = { - "filename": path / "{sector}{year}{quantity}.csv", - "quantity": "streetcred", - "sink": "nc", - } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - result = factory(config, sector_name="Yoyo")(market, None, None) - assert result[0] == path / "yoyo2010streetcred.csv" + expected_path = path / ( + f"{'Yoyo' if config_type == 'suffix' else 'yoyo'}" + f"2010" + f"{'Streetcred' if config_type == 'suffix' else 'streetcred'}" + f".{'nc' if config_type == 'suffix' else 'csv'}" + ) + assert result[0] == expected_path assert result[0].is_file() xr.open_dataset(result[0]) @mark.usefixtures("streetcred") -def test_save_with_fullpath_to_excel_with_sink(tmpdir): +def test_save_with_fullpath_to_excel(tmpdir, market): + """Test saving output to Excel file.""" from warnings import simplefilter from pandas import read_excel @@ -114,25 +135,38 @@ def test_save_with_fullpath_to_excel_with_sink(tmpdir): path = Path(tmpdir) / "results" / "stuff" / "this.xlsx" config = {"filename": path, "quantity": "streetcred", "sink": "xlsx"} - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here result = factory(config, sector_name="Yoyo")(market, None, None) assert result[0] == path - assert result[0].is_file() + assert_file_exists_and_readable(result[0]) read_excel(result[0]) -@mark.usefixtures("streetcred") -def test_no_sink_or_suffix(tmpdir): - from muse.outputs.sector import factory +@patch("muse.outputs.cache.consolidate_quantity") +def test_output_functions(mock_consolidate): + """Test output functions (capacity, production, lcoe) with common setup.""" + from muse.outputs.cache import capacity, lcoe, production + + cached = [xr.DataArray() for _ in range(3)] + agents = {} + + for func, quantity in [ + (capacity, "capacity"), + (production, "production"), + (lcoe, "lcoe"), + ]: + func(cached, agents) + mock_consolidate.assert_called_once_with(quantity, cached, agents) + mock_consolidate.reset_mock() + +@mark.usefixtures("streetcred") +def test_no_sink_or_suffix(tmpdir, market): + """Test default sink and suffix behavior.""" config = dict( quantity="streetcred", filename=f"{tmpdir}/{{Sector}}{{Quantity}}{{year}}{{suffix}}", ) - outputs = factory(config) - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - result = outputs(market, None, None) + result = factory(config)(market, None, None) assert len(result) == 1 assert result[0].is_file() assert result[0].suffix == ".csv" @@ -140,6 +174,7 @@ def test_no_sink_or_suffix(tmpdir): @mark.usefixtures("save_registries") def test_can_register_class(): + """Test class registration functionality.""" from muse.outputs.sinks import factory, register_output_sink @register_output_sink @@ -151,12 +186,14 @@ def __init__(self, sector, some_args=3): def __call__(self, x): pass + # Test default arguments settings = {"sink": {"name": "AClass"}} sink = factory(settings, sector_name="yoyo") assert isinstance(sink, AClass) assert sink.sector == "yoyo" assert sink.some_args == 3 + # Test custom arguments settings = {"sink": {"name": "AClass", "some_args": 5}} sink = factory(settings, sector_name="yoyo") assert isinstance(sink, AClass) @@ -166,6 +203,7 @@ def __call__(self, x): @mark.usefixtures("save_registries") def test_can_register_function(): + """Test function registration functionality.""" from muse.outputs.sinks import factory, register_output_sink @register_output_sink @@ -179,8 +217,10 @@ def a_function(*args): @mark.usefixtures("save_registries") def test_yearly_aggregate(): + """Test yearly aggregation with custom sink.""" from muse.outputs.sinks import factory, register_output_sink + # Setup tracking variables received_data = None gyear = None gsector = None @@ -202,16 +242,17 @@ def dummy(data, year: int, sector: str, overwrite: bool) -> MySpecialReturn: dict(overwrite=True, sink=dict(aggregate="dummy")), sector_name="yoyo" ) - data = xr.DataArray([1, 0], coords=dict(a=[2, 4]), dims="a", name="nada") + # Test first year + data = create_test_data_array([1, 0], name="nada") data["year"] = 2010 - assert isinstance(sink(data, 2010), MySpecialReturn) assert gyear == 2010 assert gsector == "yoyo" assert goverwrite is True assert isinstance(received_data, pd.DataFrame) - data = xr.DataArray([0, 1], coords=dict(a=[2, 4]), dims="a", name="nada") + # Test second year + data = create_test_data_array([0, 1], name="nada") data["year"] = 2020 assert isinstance(sink(data, 2020), MySpecialReturn) assert gyear == 2020 @@ -221,34 +262,42 @@ def dummy(data, year: int, sector: str, overwrite: bool) -> MySpecialReturn: def test_yearly_aggregate_file(tmpdir): + """Test yearly aggregation to file with multiple years of data.""" from muse.outputs.sinks import factory path = Path(tmpdir) / "file.csv" sink = factory(dict(filename=str(path), sink="aggregate"), sector_name="yoyo") - data = xr.DataArray([1, 0], coords=dict(a=[2, 4]), dims="a", name="georges") - data["year"] = 2010 - assert sink(data, 2010) == path - dataframe = pd.read_csv(path) - assert set(dataframe.columns) == {"year", "georges"} - assert dataframe.shape[0] == 2 + def verify_year_data(values, year, expected_rows): + data = create_test_data_array(values, name="georges") + data["year"] = year + assert sink(data, year) == path + df = assert_file_exists_and_readable(path, {"year", "georges"}) + assert df.shape[0] == expected_rows + return df - data = xr.DataArray([0, 1], coords=dict(a=[2, 4]), dims="a", name="georges") - data["year"] = 2020 - assert sink(data, 2020) == path - dataframe = pd.read_csv(path) - assert set(dataframe.columns) == {"year", "georges"} - assert dataframe.shape[0] == 4 + # Test first year + verify_year_data([1, 0], 2010, 2) + + # Test second year (should append to existing file) + df2 = verify_year_data([0, 1], 2020, 4) + + # Verify data from both years is present + assert set(df2.year.unique()) == {2010, 2020} + assert df2[df2.year == 2010].georges.tolist() == [1, 0] + assert df2[df2.year == 2020].georges.tolist() == [0, 1] def test_yearly_aggregate_no_outputs(tmpdir): + """Test behavior with no outputs configured.""" from muse.outputs.mca import factory outputs = factory() assert len(outputs(None, year=2010)) == 0 -def test_mca_aggregate_outputs(tmpdir): +def setup_mca_test(tmpdir, outputs_config): + """Helper function to set up MCA tests.""" from toml import dump, load from muse import examples @@ -256,17 +305,20 @@ def test_mca_aggregate_outputs(tmpdir): examples.copy_model(path=str(tmpdir)) settings = load(str(tmpdir / "model" / "settings.toml")) - settings["outputs"] = [ - dict(filename="{path}/{Quantity}{suffix}", quantity="prices", sink="aggregate") - ] + settings["outputs"] = [outputs_config] settings["time_framework"] = settings["time_framework"][:2] dump(settings, (tmpdir / "model" / "settings.toml")) + return MCA.factory(str(tmpdir / "model" / "settings.toml")) - mca = MCA.factory(str(tmpdir / "model" / "settings.toml")) - mca.run() +def test_mca_aggregate_outputs(tmpdir): + """Test MCA aggregate outputs.""" + mca = setup_mca_test( + tmpdir, + dict(filename="{path}/{Quantity}{suffix}", quantity="prices", sink="aggregate"), + ) + mca.run() assert (tmpdir / "model" / "Prices.csv").exists() - # TODO: should pass again after #612 # data = pd.read_csv(tmpdir / "model" / "Prices.csv") # assert set(data.year) == set(settings["time_framework"]) @@ -274,23 +326,10 @@ def test_mca_aggregate_outputs(tmpdir): @mark.usefixtures("save_registries") def test_path_formatting(tmpdir): - from toml import dump, load - - from muse.examples import copy_model - from muse.mca import MCA + """Test path formatting with dummy sink and quantity.""" from muse.outputs.mca import register_output_quantity from muse.outputs.sinks import register_output_sink, sink_to_file - # Copy the data to tmpdir - copy_model(path=tmpdir) - - settings_file = tmpdir / "model" / "settings.toml" - settings = load(settings_file) - settings["outputs"] = [ - dict(quantity="dummy", sink="to_dummy", filename="{path}/{Quantity}{suffix}") - ] - dump(settings, (settings_file)) - @register_output_sink(name="dummy_sink") @sink_to_file(".dummy") def to_dummy(quantity, filename, **params) -> None: @@ -300,12 +339,11 @@ def to_dummy(quantity, filename, **params) -> None: def dummy(market, **kwargs): return xr.DataArray() - mca = MCA.factory(Path(settings_file)) - assert mca.outputs(mca.market)[0] == Path( - settings["outputs"][0]["filename"].format( - path=tmpdir / "model", Quantity="Dummy", suffix=".dummy" - ) + mca = setup_mca_test( + tmpdir, + dict(quantity="dummy", sink="to_dummy", filename="{path}/{Quantity}{suffix}"), ) + assert mca.outputs(mca.market)[0] == Path(tmpdir / "model" / "Dummy.dummy") def test_register_output_quantity_cache(): @@ -319,64 +357,59 @@ def dummy_quantity(*args): class TestOutputCache: + @fixture + def output_params(self): + return [dict(quantity="height"), dict(quantity="width")] + + @fixture + def output_quantities(self, output_params): + quantities = {q["quantity"]: lambda _: None for q in output_params} + quantities["depth"] = lambda _: None + return quantities + + @fixture + def topic(self): + return "BBC Muse" + @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_init(self, mock_factory, mock_subscribe): + def test_init( + self, mock_factory, mock_subscribe, output_params, output_quantities, topic + ): from muse.outputs.cache import OutputCache - param = [dict(quantity="height"), dict(quantity="width")] - output_quantities = {q["quantity"]: lambda _: None for q in param} - output_quantities["depth"] = lambda _: None - topic = "BBC Muse" - output_cache = OutputCache( - *param, output_quantities=output_quantities, topic=topic + *output_params, output_quantities=output_quantities, topic=topic ) - - assert mock_factory.call_count == len(param) + assert mock_factory.call_count == len(output_params) mock_subscribe.assert_called_once_with(output_cache.cache, topic) @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_cache(self, mock_factory, mock_subscribe): - import xarray as xr - + def test_cache( + self, mock_factory, mock_subscribe, output_params, output_quantities, topic + ): from muse.outputs.cache import OutputCache - param = [dict(quantity="height"), dict(quantity="width")] - output_quantities = {q["quantity"]: lambda _: None for q in param} - output_quantities["depth"] = lambda _: None - topic = "BBC Muse" - output_cache = OutputCache( - *param, output_quantities=output_quantities, topic=topic + *output_params, output_quantities=output_quantities, topic=topic ) - output_cache.cache(dict(height=xr.DataArray(), depth=xr.DataArray())) - assert len(output_cache.to_save.get("height")) == 1 assert len(output_cache.to_save.get("depth", [])) == 0 @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_consolidate_cache(self, mock_factory, mock_subscribe): - import xarray as xr - + def test_consolidate_cache( + self, mock_factory, mock_subscribe, output_params, output_quantities, topic + ): from muse.outputs.cache import OutputCache - param = [dict(quantity="height"), dict(quantity="width")] - output_quantities = {q["quantity"]: lambda _: None for q in param} - output_quantities["depth"] = lambda _: None - topic = "BBC Muse" - year = 2042 - output_cache = OutputCache( - *param, output_quantities=output_quantities, topic=topic + *output_params, output_quantities=output_quantities, topic=topic ) - output_cache.cache(dict(height=xr.DataArray())) - output_cache.consolidate_cache(year) - + output_cache.consolidate_cache(2042) output_cache.factory["height"].assert_called_once() @@ -388,9 +421,12 @@ def test_cache_quantity(mock_match, mock_send): result = {"mass": 42} mock_match.return_value = result + def verify_message_sent(): + mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) + mock_send.reset_mock() + cache_quantity(**result) - mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) - mock_send.reset_mock() + verify_message_sent() with raises(ValueError): cache_quantity(function=lambda: None, mass=42) @@ -407,7 +443,7 @@ def fun2(): fun2() mock_match.assert_called_once_with("mass", 42) - mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) + verify_message_sent() def test_match_quantities(): @@ -415,37 +451,27 @@ def test_match_quantities(): from muse.outputs.cache import match_quantities - q = "mass" - da = xr.DataArray(name=q) - ds = xr.Dataset({q: da}) - def assert_equal(a: dict[str, xr.DataArray], b: dict[str, xr.DataArray]): assert set(a.keys()) == set(b.keys()) for k in a: xr.testing.assert_equal(a[k], b[k]) - actual = match_quantities(quantity=q, data=da) - assert_equal(actual, {q: da}) - - actual = match_quantities(quantity=q, data=ds) - assert_equal(actual, {q: da}) + # Test single quantity with DataArray + q = "mass" + da = xr.DataArray(name=q) + ds = xr.Dataset({q: da}) + assert_equal(match_quantities(quantity=q, data=da), {q: da}) + assert_equal(match_quantities(quantity=q, data=ds), {q: da}) + # Test multiple quantities with Dataset p = "height" ds = xr.Dataset({q: da, p: da, "rubish": da}) - actual = match_quantities(quantity=[q, p], data=ds) - assert_equal(actual, {q: da, p: da}) - - actual = match_quantities(quantity=[q, p], data=[da, da]) - assert_equal(actual, {q: da, p: da}) + assert_equal(match_quantities(quantity=[q, p], data=ds), {q: da, p: da}) + assert_equal(match_quantities(quantity=[q, p], data=[da, da]), {q: da, p: da}) + # Test error cases with raises(ValueError): - match_quantities( - quantity=[q, p], - data=[ - da, - ], - ) - + match_quantities(quantity=[q, p], data=[da]) with raises(TypeError): match_quantities(quantity=[q, p], data=42) @@ -464,51 +490,54 @@ def test_extract_agents(mock_extract): def test_extract_agents_internal(newcapa_agent, retro_agent): + """Test internal agent extraction.""" from types import SimpleNamespace from muse.outputs.cache import extract_agents_internal - newcapa_agent.name = "A1" - retro_agent.name = "A2" - sector = SimpleNamespace(name="IT", agents=[newcapa_agent, retro_agent]) + def setup_agent(agent, name): + agent.name = name + return agent + + agents = [setup_agent(newcapa_agent, "A1"), setup_agent(retro_agent, "A2")] + sector = SimpleNamespace(name="IT", agents=agents) actual = extract_agents_internal(sector) - for agent in [newcapa_agent, retro_agent]: + expected_keys = ("agent", "category", "sector", "dst_region") + + for agent in agents: assert agent.uuid in actual - assert tuple(actual[agent.uuid].keys()) == ( - "agent", - "category", - "sector", - "dst_region", - ) - assert actual[agent.uuid]["agent"] == agent.name - assert actual[agent.uuid]["category"] == agent.category - assert actual[agent.uuid]["sector"] == "IT" - assert actual[agent.uuid]["dst_region"] == agent.region + assert tuple(actual[agent.uuid].keys()) == expected_keys + agent_data = actual[agent.uuid] + assert agent_data["agent"] == agent.name + assert agent_data["category"] == agent.category + assert agent_data["sector"] == "IT" + assert agent_data["dst_region"] == agent.region def test_aggregate_cache(): import numpy as np - import xarray as xr from pandas.testing import assert_frame_equal from muse.outputs.cache import _aggregate_cache quantity = "height" - a = xr.DataArray(np.ones((3, 4, 5)), name=quantity) b = a.copy() b[0, 0, 0] = 0 + def to_df(arr): + return arr.to_dataframe().reset_index().astype(float) + actual = _aggregate_cache(quantity, [a, b]) - assert_frame_equal(actual, b.to_dataframe().reset_index().astype(float)) + assert_frame_equal(actual, to_df(b)) actual = _aggregate_cache(quantity, [b, a]) - assert_frame_equal(actual, a.to_dataframe().reset_index().astype(float)) + assert_frame_equal(actual, to_df(a)) c = a.copy() c.assign_coords(dim_0=c.dim_0.data * 10) - dc, da = (da.to_dataframe().reset_index() for da in [c, a]) + dc, da = map(to_df, [c, a]) actual = _aggregate_cache(quantity, [c, a]) expected = pd.DataFrame.merge(dc, da, how="outer").astype(float) @@ -516,81 +545,38 @@ def test_aggregate_cache(): def test_consolidate_quantity(newcapa_agent, retro_agent): + """Test consolidation of quantity data with agent information.""" from types import SimpleNamespace from muse.outputs.cache import consolidate_quantity, extract_agents_internal - newcapa_agent.name = "A1" - retro_agent.name = "A2" - newcapa_agent.category = "newcapa" - retro_agent.category = "retro" + def setup_agent(agent, name, category): + agent.name = name + agent.category = category + return agent + + newcapa_agent = setup_agent(newcapa_agent, "A1", "newcapa") + retro_agent = setup_agent(retro_agent, "A2", "retro") sector = SimpleNamespace(name="IT", agents=[newcapa_agent, retro_agent]) agents = extract_agents_internal(sector) - quantity = "height" - a = xr.DataArray( - np.ones((3, 4, 5)), - dims=("agent", "replacement", "asset"), - coords={ - "agent": [ - newcapa_agent.uuid, - ] - * 3 - }, - name=quantity, - ) - b = a.copy() - b[0, 0, 0] = 0 - b.assign_coords( - agent=[ - retro_agent.uuid, - ] - * 3 - ) + def create_agent_array(agent_uuid, modify_first=False): + arr = xr.DataArray( + np.ones((3, 4, 5)), + dims=("agent", "replacement", "asset"), + coords={"agent": [agent_uuid] * 3}, + name="height", + ) + if modify_first: + arr[0, 0, 0] = 0 + return arr - actual = consolidate_quantity(quantity, [a, b], agents) + a = create_agent_array(newcapa_agent.uuid) + b = create_agent_array(retro_agent.uuid, modify_first=True) - cols = set((*agents[retro_agent.uuid].keys(), "technology", quantity)) + actual = consolidate_quantity("height", [a, b], agents) + cols = set((*agents[retro_agent.uuid].keys(), "technology", "height")) assert set(actual.columns) == cols assert all( name in (newcapa_agent.name, retro_agent.name) for name in actual.agent.unique() ) - - -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_capacity(mock_consolidate): - import xarray as xr - - from muse.outputs.cache import capacity - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - capacity(cached, agents) - mock_consolidate.assert_called_once_with("capacity", cached, agents) - - -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_production(mock_consolidate): - import xarray as xr - - from muse.outputs.cache import production - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - production(cached, agents) - mock_consolidate.assert_called_once_with("production", cached, agents) - - -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_lcoe(mock_consolidate): - import xarray as xr - - from muse.outputs.cache import lcoe - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - lcoe(cached, agents) - mock_consolidate.assert_called_once_with("lcoe", cached, agents) From 435a7141dc9a1cae533b962c19fe6ad3bbfff239 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:20:00 +0100 Subject: [PATCH 06/33] Tidy test_timeslices --- tests/test_timeslices.py | 355 +++++++++++++++++---------------------- 1 file changed, 156 insertions(+), 199 deletions(-) diff --git a/tests/test_timeslices.py b/tests/test_timeslices.py index 283dfc63..5a31605c 100644 --- a/tests/test_timeslices.py +++ b/tests/test_timeslices.py @@ -4,126 +4,121 @@ from pytest import approx, fixture, raises from xarray import DataArray +from muse.timeslices import ( + broadcast_timeslice, + compress_timeslice, + distribute_timeslice, + drop_timeslice, + expand_timeslice, + get_level, + read_timeslices, + sort_timeslices, + timeslice_max, +) + +# Constants +TIMESLICE_LEVELS = ["month", "day", "hour"] +TIMESLICE_COORDS = {"timeslice", "month", "day", "hour"} + @fixture def non_timesliced_dataarray(): + """Create a simple non-timesliced DataArray for testing.""" return DataArray([1, 2, 3], dims=["x"]) +@fixture +def timeslice_toml(): + """TOML configuration for timeslice testing.""" + return """ + [timeslices] + winter.weekday.night = 396 + winter.weekday.morning = 396 + winter.weekday.afternoon = 264 + winter.weekend.night = 156 + winter.weekend.morning = 156 + winter.weekend.afternoon = 156 + springautumn.weekday.night = 792 + springautumn.weekday.morning = 792 + springautumn.weekday.afternoon = 528 + springautumn.weekend.night = 300 + springautumn.weekend.morning = 300 + springautumn.weekend.afternoon = 300 + summer.weekday.night = 396 + summer.weekday.morning = 396 + summer.weekday.afternoon = 264 + summer.weekend.night = 150 + summer.weekend.morning = 150 + summer.weekend.afternoon = 150 + """ + + @fixture(params=[False, True]) -def timeslice(request): - """Fixture to generate a timeslice DataArray. +def timeslice(request, timeslice_toml): + """Generate a timeslice DataArray. Creates two versions: - - A one-dimensional DataArray with a single "timeslice" dimension representing the - weights of each timeslice. - - A two-dimensional DataArray with an additional "x" dimension. This allows each "x" - coordinate to have different timeslice weights. + - A one-dimensional DataArray with a single "timeslice" dimension + - A two-dimensional DataArray with an additional "x" dimension and randomized + weights """ from toml import loads - from muse.timeslices import read_timeslices - - # Define timeslice inputs - inputs = loads( - """ - [timeslices] - winter.weekday.night = 396 - winter.weekday.morning = 396 - winter.weekday.afternoon = 264 - winter.weekend.night = 156 - winter.weekend.morning = 156 - winter.weekend.afternoon = 156 - springautumn.weekday.night = 792 - springautumn.weekday.morning = 792 - springautumn.weekday.afternoon = 528 - springautumn.weekend.night = 300 - springautumn.weekend.morning = 300 - springautumn.weekend.afternoon = 300 - summer.weekday.night = 396 - summer.weekday.morning = 396 - summer.weekday.afternoon = 264 - summer.weekend.night = 150 - summer.weekend.morning = 150 - summer.weekend.afternoon = 150 - """ - ) - - # Read timeslices - ts = read_timeslices(inputs) - assert isinstance(ts, DataArray) - assert "timeslice" in ts.coords + ts = read_timeslices(loads(timeslice_toml)) - # Expand dimensions if requested if request.param: ts = ts.expand_dims({"x": 3}) - # Add some randomization so that each "x" coord has different timeslice weights ts = ts + np.random.randint(0, 10, ts.shape) return ts def test_no_overlap(): - from pytest import raises - - from muse.timeslices import read_timeslices - + """Test that overlapping timeslice definitions raise ValueError.""" + invalid_toml = """ + [timeslices] + winter.weekday.night = 396 + winter.weekday.morning = 396 + winter.weekday.weekend = 156 + winter.weekend.night = 156 + winter.weekend.morning = 156 + winter.weekend.weekend = 156 + """ with raises(ValueError): - read_timeslices( - """ - [timeslices] - winter.weekday.night = 396 - winter.weekday.morning = 396 - winter.weekday.weekend = 156 - winter.weekend.night = 156 - winter.weekend.morning = 156 - winter.weekend.weekend = 156 - """ - ) + read_timeslices(invalid_toml) def test_drop_timeslice(non_timesliced_dataarray, timeslice): - from muse.timeslices import broadcast_timeslice, drop_timeslice + """Test dropping timeslice coordinates.""" + timesliced = broadcast_timeslice(non_timesliced_dataarray, ts=timeslice) + dropped = drop_timeslice(timesliced) - # Test on array with timeslice data - timesliced_dataarray = broadcast_timeslice(non_timesliced_dataarray, ts=timeslice) - dropped = drop_timeslice(timesliced_dataarray) - coords_to_check = {"timeslice", "month", "day", "hour"} - assert coords_to_check.issubset(timesliced_dataarray.coords) - assert not coords_to_check.intersection(dropped.coords) - - # Test on arrays without timeslice data + assert TIMESLICE_COORDS.issubset(timesliced.coords) + assert not TIMESLICE_COORDS.intersection(dropped.coords) assert drop_timeslice(non_timesliced_dataarray).equals(non_timesliced_dataarray) assert drop_timeslice(dropped).equals(dropped) def test_broadcast_timeslice(non_timesliced_dataarray, timeslice): - from muse.timeslices import broadcast_timeslice, compress_timeslice - - # Broadcast array to different levels of granularity - for level in ["month", "day", "hour"]: + """Test broadcasting arrays to different timeslice granularities.""" + for level in TIMESLICE_LEVELS: out = broadcast_timeslice(non_timesliced_dataarray, ts=timeslice, level=level) - target_timeslices = compress_timeslice( + target = compress_timeslice( timeslice, ts=timeslice, level=level, operation="sum" ) - # Check that timeslicing in output matches the global scheme - assert out.timeslice.equals(target_timeslices.timeslice) - - # Check that all timeslices in the output are equal to each other + assert out.timeslice.equals(target.timeslice) assert (out.diff(dim="timeslice") == 0).all() - - # Check that all values in the output are equal to the input assert all( (out.isel(timeslice=i) == non_timesliced_dataarray).all() for i in range(out.sizes["timeslice"]) ) - # Calling on a fully timesliced array: the input should be returned unchanged + # Test broadcasting an already timesliced array (should return unchanged) out2 = broadcast_timeslice(out, ts=timeslice) assert out2.equals(out) - # Calling on an array with inappropriate timeslicing: ValueError should be raised + # Test broadcasting with incompatible timeslice levels with raises(ValueError): broadcast_timeslice( compress_timeslice(out, ts=timeslice, level="day"), ts=timeslice @@ -131,39 +126,30 @@ def test_broadcast_timeslice(non_timesliced_dataarray, timeslice): def test_distribute_timeslice(non_timesliced_dataarray, timeslice): - from muse.timeslices import ( - broadcast_timeslice, - compress_timeslice, - distribute_timeslice, - ) - - # Distribute array to different levels of granularity - for level in ["month", "day", "hour"]: + """Test distributing arrays across timeslices.""" + for level in TIMESLICE_LEVELS: out = distribute_timeslice(non_timesliced_dataarray, ts=timeslice, level=level) - target_timeslices = compress_timeslice( + target = compress_timeslice( timeslice, ts=timeslice, level=level, operation="sum" ) - # Check that timeslicing in output matches the global scheme - assert out.timeslice.equals(target_timeslices.timeslice) + assert out.timeslice.equals(target.timeslice) - # Check that all values are proportional to timeslice lengths - out_proportions = out / broadcast_timeslice( + # Check proportionality + out_prop = out / broadcast_timeslice( out.sum("timeslice"), ts=timeslice, level=level ) - ts_proportions = target_timeslices / broadcast_timeslice( - target_timeslices.sum("timeslice"), ts=timeslice, level=level + ts_prop = target / broadcast_timeslice( + target.sum("timeslice"), ts=timeslice, level=level ) - assert abs(out_proportions - ts_proportions).max() < 1e-6 + assert abs(out_prop - ts_prop).max() < 1e-6 - # Check that the sum across timeslices is equal to the input assert (out.sum("timeslice") == approx(non_timesliced_dataarray)).all() - # Calling on a fully timesliced array: the input should be returned unchanged - out2 = distribute_timeslice(out, ts=timeslice) - assert out2.equals(out) + # Test distributing an already timesliced array (should return unchanged) + assert distribute_timeslice(out, ts=timeslice).equals(out) - # Calling on an array with inappropraite timeslicing: ValueError should be raised + # Test distributing with incompatible timeslice levels with raises(ValueError): distribute_timeslice( compress_timeslice(out, ts=timeslice, level="day"), ts=timeslice @@ -171,135 +157,106 @@ def test_distribute_timeslice(non_timesliced_dataarray, timeslice): def test_compress_timeslice(non_timesliced_dataarray, timeslice): - from muse.timeslices import broadcast_timeslice, compress_timeslice, get_level - - # Create timesliced dataarray for compressing - timesliced_dataarray = broadcast_timeslice(non_timesliced_dataarray, ts=timeslice) - - # Compress array to different levels of granularity - for level in ["month", "day", "hour"]: - # Sum operation - out = compress_timeslice( - timesliced_dataarray, ts=timeslice, operation="sum", level=level - ) - assert get_level(out) == level - assert ( - out.sum("timeslice") == approx(timesliced_dataarray.sum("timeslice")) - ).all() - - # Mean operation - out = compress_timeslice( - timesliced_dataarray, ts=timeslice, operation="mean", level=level - ) - assert get_level(out) == level - assert ( - out.mean("timeslice") == approx(timesliced_dataarray.mean("timeslice")) - ).all() # NB in general this should be a weighted mean, but this works here - # because the data is equal in every timeslice - - # Calling without specifying a level: the input should be returned unchanged - out = compress_timeslice(timesliced_dataarray, ts=timeslice) - assert out.equals(timesliced_dataarray) - - # Calling with an invalid level: ValueError should be raised + """Test compressing timesliced arrays.""" + timesliced = broadcast_timeslice(non_timesliced_dataarray, ts=timeslice) + + for level in TIMESLICE_LEVELS: + for operation in ["sum", "mean"]: + out = compress_timeslice( + timesliced, ts=timeslice, operation=operation, level=level + ) + assert get_level(out) == level + + if operation == "sum": + assert ( + out.sum("timeslice") == approx(timesliced.sum("timeslice")) + ).all() + else: # mean + assert ( + out.mean("timeslice") == approx(timesliced.mean("timeslice")) + ).all() + + # Test compressing without specifying a level (should return unchanged) + assert compress_timeslice(timesliced, ts=timeslice).equals(timesliced) + + # Test compressing with invalid level name with raises(ValueError): - compress_timeslice(timesliced_dataarray, ts=timeslice, level="invalid") + compress_timeslice(timesliced, ts=timeslice, level="invalid") - # Calling with an invalid operation: ValueError should be raised + # Test compressing with invalid operation type with raises(ValueError): - compress_timeslice( - timesliced_dataarray, ts=timeslice, level="day", operation="invalid" - ) + compress_timeslice(timesliced, ts=timeslice, level="day", operation="invalid") def test_expand_timeslice(non_timesliced_dataarray, timeslice): - from muse.timeslices import broadcast_timeslice, expand_timeslice - - # Different starting points for expansion - for level in ["month", "day", "hour"]: - timesliced_dataarray = broadcast_timeslice( + """Test expanding timesliced arrays.""" + for level in TIMESLICE_LEVELS: + timesliced = broadcast_timeslice( non_timesliced_dataarray, ts=timeslice, level=level ) - # Broadcast operation - out = expand_timeslice( - timesliced_dataarray, ts=timeslice, operation="broadcast" - ) - assert out.timeslice.equals(timeslice.timeslice) - assert ( - out.mean("timeslice") == approx(timesliced_dataarray.mean("timeslice")) - ).all() - - # Distribute operation - out = expand_timeslice( - timesliced_dataarray, ts=timeslice, operation="distribute" - ) - assert out.timeslice.equals(timeslice.timeslice) - assert ( - out.sum("timeslice") == approx(timesliced_dataarray.sum("timeslice")) - ).all() + for operation in ["broadcast", "distribute"]: + out = expand_timeslice(timesliced, ts=timeslice, operation=operation) + assert out.timeslice.equals(timeslice.timeslice) - # Calling on an already expanded array: the input should be returned unchanged - out2 = expand_timeslice(out, ts=timeslice) - assert out.equals(out2) + if operation == "broadcast": + assert ( + out.mean("timeslice") == approx(timesliced.mean("timeslice")) + ).all() + else: # distribute + assert ( + out.sum("timeslice") == approx(timesliced.sum("timeslice")) + ).all() - # Calling with an invalid operation: ValueError should be raised + # Test expanding an already expanded array (should return unchanged) + assert expand_timeslice(out, ts=timeslice).equals(out) + + # Test expanding with invalid operation type + timesliced = broadcast_timeslice( + non_timesliced_dataarray, ts=timeslice, level="month" + ) with raises(ValueError): - timesliced_dataarray = broadcast_timeslice( - non_timesliced_dataarray, ts=timeslice, level="month" - ) - expand_timeslice(timesliced_dataarray, ts=timeslice, operation="invalid") + expand_timeslice(timesliced, ts=timeslice, operation="invalid") def test_get_level(non_timesliced_dataarray, timeslice): - from muse.timeslices import broadcast_timeslice, get_level - - for level in ["month", "day", "hour"]: - timesliced_dataarray = broadcast_timeslice( + """Test getting timeslice level.""" + for level in TIMESLICE_LEVELS: + timesliced = broadcast_timeslice( non_timesliced_dataarray, ts=timeslice, level=level ) - assert get_level(timesliced_dataarray) == level + assert get_level(timesliced) == level - # Should raise error with non-timesliced array with raises(ValueError): get_level(non_timesliced_dataarray) def test_sort_timeslices(non_timesliced_dataarray, timeslice): - from muse.timeslices import broadcast_timeslice, sort_timeslices - - # Finest timeslice level -> should match ordering of `timeslice` - timesliced_dataarray = broadcast_timeslice( + """Test sorting timeslices.""" + # Test hour level + timesliced = broadcast_timeslice( non_timesliced_dataarray, ts=timeslice, level="hour" ) - sorted = sort_timeslices(timesliced_dataarray, timeslice) - assert sorted.timeslice.equals(timeslice.timeslice) - assert not sorted.timeslice.equals( - timesliced_dataarray.sortby("timeslice").timeslice - ) # but could be true if the timeslices in `timeslice` are in alphabetical order - - # Coarser timeslice level -> should match xarray sortby - timesliced_dataarray = broadcast_timeslice( + sorted_data = sort_timeslices(timesliced, timeslice) + assert sorted_data.timeslice.equals(timeslice.timeslice) + + # Test month level + timesliced = broadcast_timeslice( non_timesliced_dataarray, ts=timeslice, level="month" ) - sorted = sort_timeslices(timesliced_dataarray, timeslice) - assert sorted.timeslice.equals(timesliced_dataarray.sortby("timeslice").timeslice) + sorted_data = sort_timeslices(timesliced, timeslice) + assert sorted_data.timeslice.equals(timesliced.sortby("timeslice").timeslice) def test_timeslice_max(non_timesliced_dataarray): - from muse.timeslices import broadcast_timeslice, read_timeslices, timeslice_max - - # With two equal timeslice lengths, this should be equivalent to max * 2 - ts = read_timeslices( - """ - [timeslices] - winter.weekday.night = 396 - winter.weekday.morning = 396 - """ - ) - timesliced_dataarray = broadcast_timeslice(non_timesliced_dataarray, ts=ts) - timesliced_dataarray = timesliced_dataarray + np.random.rand( - *timesliced_dataarray.shape - ) - timeslice_max_dataarray = timeslice_max(timesliced_dataarray, ts=ts) - assert timeslice_max_dataarray.equals(timesliced_dataarray.max("timeslice") * 2) + """Test timeslice maximum calculation.""" + ts = read_timeslices(""" + [timeslices] + winter.weekday.night = 396 + winter.weekday.morning = 396 + """) + + timesliced = broadcast_timeslice(non_timesliced_dataarray, ts=ts) + timesliced = timesliced + np.random.rand(*timesliced.shape) + max_val = timeslice_max(timesliced, ts=ts) + assert max_val.equals(timesliced.max("timeslice") * 2) From cc3310d31d9cb543733326ae4fb07d74f1f6bf99 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:25:38 +0100 Subject: [PATCH 07/33] Tidy test_trade --- tests/test_trade.py | 72 ++++++++++++++++++++++----------------------- 1 file changed, 35 insertions(+), 37 deletions(-) diff --git a/tests/test_trade.py b/tests/test_trade.py index 14f51c73..8191a09d 100644 --- a/tests/test_trade.py +++ b/tests/test_trade.py @@ -3,17 +3,19 @@ from pytest import approx, fixture +from muse import constraints as cs +from muse import examples +from muse.utilities import agent_concatenation, reduce_assets + @fixture def constraints_args(sector="power", model="trade") -> Mapping[str, Any]: - from muse import examples - from muse.utilities import agent_concatenation, reduce_assets - power = examples.sector(model=model, sector=sector) search_space = examples.search_space("power", model="trade") market = examples.matching_market("power", "trade") - assets = agent_concatenation({u.uuid: u.assets for u in list(power.agents)}) + assets = agent_concatenation({agent.uuid: agent.assets for agent in power.agents}) capacity = reduce_assets(assets.capacity, coords=("region", "technology")) + return dict( demand=market.consumption.sel(year=2025, drop=True), capacity=capacity.sel(year=[2020, 2025]), @@ -24,59 +26,48 @@ def constraints_args(sector="power", model="trade") -> Mapping[str, Any]: def test_demand_constraint(constraints_args): - from muse import constraints as cs - constraint = cs.demand(**constraints_args) assert set(constraint.b.dims) == {"timeslice", "dst_region", "commodity"} def test_max_capacity_constraints(constraints_args): - from muse import constraints as cs - constraint = cs.max_capacity_expansion(**constraints_args) assert constraint.production == 0 assert set(constraint.capacity.dims) == {"agent", "src_region"} - assert ((constraint.region == constraint.src_region) == constraint.capacity).all() assert set(constraint.b.dims) == {"replacement", "dst_region", "src_region"} assert set(constraint.agent.coords) == {"region", "agent"} + assert ((constraint.region == constraint.src_region) == constraint.capacity).all() def test_max_production(constraints_args): - from muse import constraints as cs - constraint = cs.max_production(**constraints_args) - dims = { + production_dims = { "timeslice", "commodity", "replacement", "agent", - "timeslice", "dst_region", "src_region", } - assert set(constraint.capacity.dims) == dims - assert set(constraint.production.dims) == dims + assert set(constraint.capacity.dims) == production_dims + assert set(constraint.production.dims) == production_dims assert set(constraint.agent.coords) == {"region", "agent"} def test_minimum_service(constraints_args): - from muse import constraints as cs - assert cs.minimum_service(**constraints_args) is None constraints_args["technologies"]["minimum_service_factor"] = 0.5 constraint = cs.minimum_service(**constraints_args) - dims = {"replacement", "agent", "commodity", "timeslice"} - assert set(constraint.capacity.dims) == dims - assert set(constraint.production.dims) == dims - assert set(constraint.b.dims) == dims - assert (constraint.capacity <= 0).all() + service_dims = {"replacement", "agent", "commodity", "timeslice"} + assert set(constraint.capacity.dims) == service_dims + assert set(constraint.production.dims) == service_dims + assert set(constraint.b.dims) == service_dims assert set(constraint.agent.coords) == {"region", "agent"} + assert (constraint.capacity <= 0).all() def test_search_space(constraints_args): - from muse import constraints as cs - search_space = constraints_args["search_space"] search_space[:] = 1 assert cs.search_space(**constraints_args) is None @@ -90,31 +81,38 @@ def test_search_space(constraints_args): assert set(constraint.agent.coords) == {"region", "agent"} -def test_power_sector_no_investment(): - from muse import examples - from muse.utilities import agent_concatenation +def get_agent_capacities(sector): + """Helper to get concatenated agent capacities.""" + return agent_concatenation( + {agent.uuid: agent.assets.capacity for agent in sector.agents} + ) + +def test_power_sector_no_investment(): power = examples.sector("power", "trade") market = examples.matching_market("power", "trade").sel(year=[2020, 2025]) - initial = agent_concatenation({u.uuid: u.assets.capacity for u in power.agents}) + initial_capacity = get_agent_capacities(power) power.next(market) - final = agent_concatenation({u.uuid: u.assets.capacity for u in power.agents}) + final_capacity = get_agent_capacities(power) - assert (initial == final).all() + assert (initial_capacity == final_capacity).all() def test_power_sector_some_investment(): - from muse import examples - from muse.utilities import agent_concatenation - power = examples.sector("power", "trade") market = examples.matching_market("power", "trade").sel(year=[2020, 2025]) market.consumption[:] *= 1.5 - initial = agent_concatenation({u.uuid: u.assets.capacity for u in power.agents}) + initial_capacity = get_agent_capacities(power) result = power.next(market) - final = agent_concatenation({u.uuid: u.assets.capacity for u in power.agents}) - assert "windturbine" not in initial.technology - assert final.sel(asset=final.technology == "windturbine", year=2025).sum() < 1 + final_capacity = get_agent_capacities(power) + + assert "windturbine" not in initial_capacity.technology + assert ( + final_capacity.sel( + asset=final_capacity.technology == "windturbine", year=2025 + ).sum() + < 1 + ) assert "dst_region" not in result.dims From d9f41835b4f1bc726d119c07025b8721953876e9 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:28:22 +0100 Subject: [PATCH 08/33] Tidy test_timeslice_output --- tests/test_timeslice_output.py | 176 ++++++++++++--------------------- 1 file changed, 64 insertions(+), 112 deletions(-) diff --git a/tests/test_timeslice_output.py b/tests/test_timeslice_output.py index 6212a81e..4538a744 100644 --- a/tests/test_timeslice_output.py +++ b/tests/test_timeslice_output.py @@ -1,75 +1,68 @@ +from operator import ge, le +from pathlib import Path + +import pandas as pd from pytest import mark, raises +from muse import examples +from muse.mca import MCA -def modify_technodata_timeslices(model_path, sector, process_name, utilization_factors): - import pandas as pd - technodata_timeslices = pd.read_csv( - model_path / sector / "TechnodataTimeslices.csv" +def setup_test_environment( + tmpdir: Path, + sector: str, + process_names: tuple[str, str], + utilization_factors: list[float], +) -> Path: + """Set up test environment with modified technodata timeslices.""" + model_path = examples.copy_model( + name="default_timeslice", path=tmpdir, overwrite=True ) - technodata_timeslices.loc[ - technodata_timeslices["ProcessName"] == process_name[0], "UtilizationFactor" - ] = utilization_factors[0] + # Read and modify technodata timeslices + technodata = pd.read_csv(model_path / sector / "TechnodataTimeslices.csv") + for process, factor in zip(process_names, utilization_factors): + technodata.loc[technodata["ProcessName"] == process, "UtilizationFactor"] = ( + factor + ) + technodata["MinimumServiceFactor"] = 0 - technodata_timeslices.loc[ - technodata_timeslices["ProcessName"] == process_name[1], "UtilizationFactor" - ] = utilization_factors[1] - technodata_timeslices["MinimumServiceFactor"] = 0 - return technodata_timeslices + # Save modified data + output_path = model_path / sector / "TechnodataTimeslices.csv" + technodata.to_csv(output_path, index=False) + return model_path -@mark.parametrize("utilization_factors", [([0.1], [1]), ([1], [0.1])]) -@mark.parametrize("process_name", [("gasCCGT", "windturbine")]) -def test_fullsim_timeslices(tmpdir, utilization_factors, process_name): - from operator import ge, le +PROCESS_PAIR = [("gasCCGT", "windturbine")] - import pandas as pd - - from muse import examples - from muse.mca import MCA +@mark.parametrize("utilization_factors", [([0.1], [1]), ([1], [0.1])]) +@mark.parametrize("process_names", PROCESS_PAIR) +def test_fullsim_timeslices(tmpdir, utilization_factors, process_names): sector = "power" - - # Copy the model inputs to tmpdir - model_path = examples.copy_model( - name="default_timeslice", path=tmpdir, overwrite=True - ) - technodata_timeslices = modify_technodata_timeslices( - model_path=model_path, - sector=sector, - process_name=process_name, - utilization_factors=utilization_factors, - ) - - technodata_timeslices.to_csv( - model_path / sector / "TechnodataTimeslices.csv", index=False + model_path = setup_test_environment( + tmpdir, sector, process_names, utilization_factors ) with tmpdir.as_cwd(): MCA.factory(model_path / "settings.toml").run() - MCACapacity = pd.read_csv(tmpdir / "Results/MCACapacity.csv") - - if utilization_factors[0] > utilization_factors[1]: - operator = ge - else: - operator = le - - assert operator( - len( - MCACapacity[ - (MCACapacity.sector == sector) - & (MCACapacity.technology == process_name[0]) - ] - ), - len( - MCACapacity[ - (MCACapacity.sector == sector) - & (MCACapacity.technology == process_name[1]) - ] - ), + mca_capacity = pd.read_csv(tmpdir / "Results/MCACapacity.csv") + operator = ge if utilization_factors[0] > utilization_factors[1] else le + + tech1_count = len( + mca_capacity[ + (mca_capacity.sector == sector) + & (mca_capacity.technology == process_names[0]) + ] ) + tech2_count = len( + mca_capacity[ + (mca_capacity.sector == sector) + & (mca_capacity.technology == process_names[1]) + ] + ) + assert operator(tech1_count, tech2_count) @mark.parametrize( @@ -79,77 +72,36 @@ def test_fullsim_timeslices(tmpdir, utilization_factors, process_name): ([1, 1, 1, 1, 1, 1], [1, 1, 0.0001, 0.0001, 1, 1]), ], ) -@mark.parametrize("process_name", [("gasCCGT", "windturbine")]) +@mark.parametrize("process_names", PROCESS_PAIR) def test_zero_utilization_factor_supply_timeslice( - tmpdir, utilization_factors, process_name + tmpdir, utilization_factors, process_names ): - import pandas as pd - - from muse import examples - from muse.mca import MCA - sector = "power" - - # Copy the model inputs to tmpdir - model_path = examples.copy_model( - name="default_timeslice", path=tmpdir, overwrite=True - ) - - technodata_timeslices = modify_technodata_timeslices( - model_path=model_path, - sector=sector, - process_name=process_name, - utilization_factors=utilization_factors, - ) - - technodata_timeslices.to_csv( - model_path / sector / "TechnodataTimeslices.csv", index=False + model_path = setup_test_environment( + tmpdir, sector, process_names, utilization_factors ) with tmpdir.as_cwd(): MCA.factory(model_path / "settings.toml").run() - path = str(tmpdir / "Results" / "Power_Supply.csv") - - output = pd.read_csv(path) + power_supply = pd.read_csv(tmpdir / "Results/Power_Supply.csv").reset_index() + zero_utilization_indices = [ + i for i, factor in enumerate(utilization_factors) if factor == 0 + ] - output = output.reset_index() - zero_utilization_factors = [i for i, e in enumerate(utilization_factors) if e == 0] - - assert ( - len( - output[ - ( - output.timeslice.isin(zero_utilization_factors) - & (output.technology == process_name) - ) - ] - ) - == 0 - ) + zero_output = power_supply[ + power_supply.timeslice.isin(zero_utilization_indices) + & (power_supply.technology == process_names) + ] + assert len(zero_output) == 0 @mark.parametrize("utilization_factors", [([0], [1]), ([1], [0])]) -@mark.parametrize("process_name", [("gasCCGT", "windturbine")]) -def test_all_zero_fatal_error(tmpdir, utilization_factors, process_name): - from muse import examples - from muse.mca import MCA - +@mark.parametrize("process_names", PROCESS_PAIR) +def test_all_zero_fatal_error(tmpdir, utilization_factors, process_names): sector = "power" - - # Copy the model inputs to tmpdir - model_path = examples.copy_model( - name="default_timeslice", path=tmpdir, overwrite=True - ) - technodata_timeslices = modify_technodata_timeslices( - model_path=model_path, - sector=sector, - process_name=process_name, - utilization_factors=utilization_factors, - ) - - technodata_timeslices.to_csv( - model_path / sector / "TechnodataTimeslices.csv", index=False + model_path = setup_test_environment( + tmpdir, sector, process_names, utilization_factors ) with tmpdir.as_cwd(), raises(ValueError): From 724decc39d64689d01dfb46bd06ba7d8718bb6bb Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:29:13 +0100 Subject: [PATCH 09/33] Tidy test_quantities --- tests/test_quantities.py | 203 +++++++++++++++++---------------------- 1 file changed, 86 insertions(+), 117 deletions(-) diff --git a/tests/test_quantities.py b/tests/test_quantities.py index 1877f47e..7329ce46 100644 --- a/tests/test_quantities.py +++ b/tests/test_quantities.py @@ -2,11 +2,24 @@ import xarray as xr from pytest import approx, fixture, mark +from muse.commodities import is_enduse, is_pollutant +from muse.quantities import ( + capacity_in_use, + consumption, + emission, + maximum_production, + minimum_production, + production_amplitude, + supply, +) +from muse.timeslices import broadcast_timeslice, distribute_timeslice +from muse.utilities import broadcast_over_assets -@fixture -def technologies(technologies, capacity, timeslice): - from muse.utilities import broadcast_over_assets +@fixture +def technologies( + technologies: xr.Dataset, capacity: xr.DataArray, timeslice +) -> xr.Dataset: return broadcast_over_assets(technologies, capacity) @@ -14,8 +27,6 @@ def technologies(technologies, capacity, timeslice): def production( technologies: xr.Dataset, capacity: xr.DataArray, timeslice ) -> xr.DataArray: - from muse.timeslices import broadcast_timeslice, distribute_timeslice - return ( broadcast_timeslice(capacity) * distribute_timeslice(technologies.fixed_outputs) @@ -23,19 +34,17 @@ def production( ) -def test_consumption(technologies, production, market): - from muse.quantities import consumption - - # Prices not provided, so flexible inputs are ignored +def test_consumption(technologies: xr.Dataset, production: xr.DataArray, market): + # Test without prices consump = consumption(technologies, production) assert set(production.dims) == set(consump.dims) - # Prices provided, but no flexible inputs -> should be the same as above + # Test with prices but no flexible inputs technologies.flexible_inputs[:] = 0 consump2 = consumption(technologies, production, market.prices) assert consump2.values == approx(consump.values) - # Flexible inputs considered + # Test with prices and flexible inputs consump3 = consumption(technologies, production, market.prices) assert set(production.dims) == set(consump3.dims) @@ -43,60 +52,35 @@ def test_consumption(technologies, production, market): def test_production_aggregate_asset_view( technologies: xr.Dataset, capacity: xr.DataArray ): - """Production when capacity has format of agent.sector. + """Test production when capacity has format of agent.sector.""" + technologies = technologies[["fixed_outputs", "utilization_factor"]] + enduses = is_enduse(technologies.comm_usage) + assert enduses.any() - E.g. capacity aggregated across agents. - """ - from muse.commodities import is_enduse - from muse.quantities import maximum_production + def check_production(fouts: float, ufact: float): + technologies.fixed_outputs[:] = fouts + technologies.utilization_factor[:] = ufact + prod = maximum_production(technologies, capacity) - technologies: xr.Dataset = technologies[ # type:ignore - ["fixed_outputs", "utilization_factor"] - ] + assert set(prod.dims) == set(capacity.dims).union({"commodity", "timeslice"}) + assert prod.sel(commodity=~enduses).values == approx(0) - enduses = is_enduse(technologies.comm_usage) - assert enduses.any() + prod, expected = xr.broadcast( + prod.sel(commodity=enduses).sum("timeslice"), capacity + ) + assert prod.values == approx(fouts * ufact * expected.values) - technologies.fixed_outputs[:] = 1 - technologies.utilization_factor[:] = 1 - prod = maximum_production(technologies, capacity) - assert set(prod.dims) == set(capacity.dims).union({"commodity", "timeslice"}) - assert prod.sel(commodity=~enduses).values == approx(0) - prod, expected = xr.broadcast( - prod.sel(commodity=enduses).sum("timeslice"), capacity - ) - assert prod.values == approx(expected.values) - - technologies.fixed_outputs[:] = fouts = 2 - technologies.utilization_factor[:] = ufact = 0.5 - prod = maximum_production(technologies, capacity) - assert prod.sel(commodity=~enduses).values == approx(0) - assert set(prod.dims) == set(capacity.dims).union({"commodity", "timeslice"}) - prod, expected = xr.broadcast( - prod.sel(commodity=enduses).sum("timeslice"), capacity - ) - assert prod.values == approx(fouts * ufact * expected.values) - - technologies.fixed_outputs[:] = fouts = 3 - technologies.utilization_factor[:] = ufact = 0.5 - prod = maximum_production(technologies, capacity) - assert prod.sel(commodity=~enduses).values == approx(0) - assert set(prod.dims) == set(capacity.dims).union({"commodity", "timeslice"}) - prod, expected = xr.broadcast( - prod.sel(commodity=enduses).sum("timeslice"), capacity - ) - assert prod.values == approx(fouts * ufact * expected.values) + # Test different combinations of fixed outputs and utilization factors + check_production(1.0, 1.0) + check_production(2.0, 0.5) + check_production(3.0, 0.5) @mark.xfail def test_production_agent_asset_view( technologies: xr.Dataset, capacity: xr.DataArray, timeslice ): - """Production when capacity has format of agent.assets.capacity. - - TODO: This requires a fully-explicit technologies dataset. Need to rework the - fixtures. - """ + """Test production when capacity has format of agent.assets.capacity.""" from muse.utilities import coords_to_multiindex, reduce_assets capacity = coords_to_multiindex(reduce_assets(capacity)).unstack("asset").fillna(0) @@ -104,26 +88,25 @@ def test_production_agent_asset_view( def test_capacity_in_use(production: xr.DataArray, technologies: xr.Dataset): - from muse.commodities import is_enduse - from muse.quantities import capacity_in_use - - technologies: xr.Dataset = technologies[ # type: ignore - ["fixed_outputs", "utilization_factor"] - ] + technologies = technologies[["fixed_outputs", "utilization_factor"]] production[:] = prod = 10 technologies.fixed_outputs[:] = fout = 5 technologies.utilization_factor[:] = ufac = 2 enduses = is_enduse(technologies.comm_usage) + + # Test with max_dim=None capa = capacity_in_use(production, technologies, max_dim=None) assert "commodity" in capa.dims capa, expected = xr.broadcast(capa, enduses * prod / fout / ufac) assert capa.values == approx(expected.values) + # Test without max_dim capa = capacity_in_use(production, technologies) assert "commodity" not in capa.dims assert capa.values == approx(prod / fout / ufac) + # Test with modified production for specific commodity maxcomm = np.random.choice(production.commodity.sel(commodity=enduses).values) production.loc[{"commodity": maxcomm}] = prod = 11 capa = capacity_in_use(production, technologies) @@ -132,110 +115,96 @@ def test_capacity_in_use(production: xr.DataArray, technologies: xr.Dataset): def test_emission(production: xr.DataArray, technologies: xr.Dataset): - from muse.commodities import is_pollutant - from muse.quantities import emission - em = emission(production, technologies) - - # Check that all environmental commodities are in the result envs = is_pollutant(technologies.comm_usage) - assert em.commodity.isin(envs.commodity).all() - # Check that no non-environmental commodities are in the result + # Check environmental commodities + assert em.commodity.isin(envs.commodity).all() assert set(em.commodity.values) == set(envs.commodity[envs].values) - # If fixed_outputs for env commodities are zero, then emissions should be zero + # Test zero emissions cases techs = technologies.copy() techs.fixed_outputs.loc[{"commodity": envs}] = 0 - em = emission(production, techs) + em_zero = emission(production, techs) + # Check that all non-NaN values are zero + assert (em_zero.where(~np.isnan(em_zero), 0) == 0).all() - # If production is zero, then emissions should be zero - em = emission(production * 0, technologies) - assert (em == 0).all() + # Test zero production case + em_zero_prod = emission(production * 0, technologies) + assert (em_zero_prod.where(~np.isnan(em_zero_prod), 0) == 0).all() -def test_min_production(technologies, capacity, timeslice): +def test_min_production(technologies: xr.Dataset, capacity: xr.DataArray, timeslice): """Test minimum production quantity.""" - from muse.quantities import maximum_production, minimum_production - - # If no minimum service factor is defined, the minimum production is zero + # Test without minimum service factor assert "minimum_service_factor" not in technologies - production = minimum_production(technologies, capacity) - assert (production == 0).all() + assert (minimum_production(technologies, capacity) == 0).all() - # If minimum service factor is defined, then the minimum production is not zero - # and it is less than the maximum production + # Test with minimum service factor technologies["minimum_service_factor"] = 0.5 production = minimum_production(technologies, capacity) assert not (production == 0).all() assert (production <= maximum_production(technologies, capacity)).all() -def test_supply_single_region(technologies, capacity, production, timeslice): - from muse.commodities import is_enduse - from muse.quantities import supply - - # Select data for a single region +def test_supply_single_region( + technologies: xr.Dataset, + capacity: xr.DataArray, + production: xr.DataArray, + timeslice, +): + # Setup single region data region = "USA" technologies = technologies.where(technologies.region == region, drop=True) capacity = capacity.where(capacity.region == region, drop=True) production = production.where(production.region == region, drop=True) - # Random demand within the bounds of the maximum production - demand = production.sum("asset") - demand = demand * np.random.rand(*demand.shape) + # Create random demand + demand = production.sum("asset") * np.random.rand(*production.sum("asset").shape) assert "region" not in demand.dims - # Calculate supply - spl = supply(capacity, demand, technologies) - - # Total supply across assets should equal demand (for end-use commodities) - spl = spl.sum("asset") + # Test supply matches demand for end-use commodities + spl = supply(capacity, demand, technologies).sum("asset") enduses = is_enduse(technologies.comm_usage) assert abs(spl.sel(commodity=enduses) - demand.sel(commodity=enduses)).sum() < 1e-5 -def test_supply_multi_region(technologies, capacity, production, timeslice): - from muse.commodities import is_enduse - from muse.quantities import supply - - # Random demand within the bounds of the maximum production +def test_supply_multi_region( + technologies: xr.Dataset, + capacity: xr.DataArray, + production: xr.DataArray, + timeslice, +): + # Create random demand demand = production.groupby("region").sum("asset") demand = demand * np.random.rand(*demand.shape) - - # Calculate supply assert "region" in demand.dims - spl = supply(capacity, demand, technologies) - # Total supply across assets within each region should equal demand - # (for end-use commodities) - spl = spl.groupby("region").sum("asset") + # Test supply matches demand for end-use commodities by region + spl = supply(capacity, demand, technologies).groupby("region").sum("asset") enduses = is_enduse(technologies.comm_usage) assert abs(spl.sel(commodity=enduses) - demand.sel(commodity=enduses)).sum() < 1e-5 -def test_supply_with_min_service(technologies, capacity, production, timeslice): - from muse.quantities import minimum_production, supply - - # Calculate minimum production +def test_supply_with_min_service( + technologies: xr.Dataset, + capacity: xr.DataArray, + production: xr.DataArray, + timeslice, +): technologies["minimum_service_factor"] = 0.3 minprod = minimum_production(technologies, capacity) - # Random demand within the bounds of the maximum production + # Create random demand demand = production.groupby("region").sum("asset") demand = demand * np.random.rand(*demand.shape) - # Calculate supply + # Test supply meets minimum production constraint spl = supply(capacity, demand, technologies) - - # Supply should be greater than or equal to the minimum production assert (spl >= minprod).all() -def test_production_amplitude(production, technologies): - from muse.quantities import production_amplitude - from muse.utilities import broadcast_over_assets - +def test_production_amplitude(production: xr.DataArray, technologies: xr.Dataset): techs = broadcast_over_assets(technologies, production) result = production_amplitude(production, techs) assert set(result.dims) == set(production.dims) - {"commodity"} From 6b4c7044794517a9803fe708f17eefc8f397d476 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:29:52 +0100 Subject: [PATCH 10/33] Tidy test_mca --- tests/test_mca.py | 236 ++++++++++++++++++++++++---------------------- 1 file changed, 121 insertions(+), 115 deletions(-) diff --git a/tests/test_mca.py b/tests/test_mca.py index 63237434..0adca94f 100644 --- a/tests/test_mca.py +++ b/tests/test_mca.py @@ -1,67 +1,71 @@ from collections.abc import Sequence - -from xarray import Dataset - -from muse.commodities import CommodityUsage +from copy import deepcopy +from unittest.mock import patch + +import numpy as np +from pytest import approx +from xarray import DataArray, Dataset, broadcast + +from muse.commodities import ( + CommodityUsage, + is_consumable, + is_enduse, + is_other, +) +from muse.mca import check_demand_fulfillment, check_equilibrium, find_equilibrium from muse.timeslices import drop_timeslice def test_check_equilibrium(market: Dataset): - """Test for the equilibrium function of the MCA.""" - from muse.mca import check_equilibrium - + """Test market equilibrium checking for both demand and prices.""" years = [2010, 2020] tol = 0.1 - equilibrium_variable = "demand" - market = market.interp(year=years) new_market = market.copy(deep=True) - assert check_equilibrium(new_market, market, tol, equilibrium_variable) + # Test demand equilibrium + assert check_equilibrium(new_market, market, tol, "demand") new_market["supply"] = drop_timeslice(new_market["supply"]) + tol * 1.5 + assert not check_equilibrium(new_market, market, tol, "demand") - assert not check_equilibrium(new_market, market, tol, equilibrium_variable) - - equilibrium_variable = "prices" - - assert check_equilibrium(new_market, market, tol, equilibrium_variable) + # Test price equilibrium + assert check_equilibrium(new_market, market, tol, "prices") new_market["prices"] = drop_timeslice(new_market["prices"]) + tol * 1.5 - assert not check_equilibrium(new_market, market, tol, equilibrium_variable) + assert not check_equilibrium(new_market, market, tol, "prices") -def test_check_demand_fulfillment(market): - """Test for the demand fulfillment function of the MCA.""" - from muse.mca import check_demand_fulfillment +def test_check_demand_fulfillment(market: Dataset): + """Test if market demand is fulfilled within tolerance.""" + tolerance = -0.1 + market["supply"] = drop_timeslice(market.consumption.copy(deep=True)) - tolerance_unmet_demand = -0.1 + assert check_demand_fulfillment(market, tolerance) - market["supply"] = drop_timeslice(market.consumption.copy(deep=True)) - assert check_demand_fulfillment( - market, - tolerance_unmet_demand, - ) - market["supply"] = drop_timeslice(market["supply"]) + tolerance_unmet_demand * 1.5 - assert not check_demand_fulfillment( - market, - tolerance_unmet_demand, - ) + # Test with unfulfilled demand + market["supply"] = drop_timeslice(market["supply"]) + tolerance * 1.5 + assert not check_demand_fulfillment(market, tolerance) def sector_market(market: Dataset, comm_usage: Sequence[CommodityUsage]) -> Dataset: - """Creates a likely return market from a sector.""" - from numpy.random import randint - from xarray import DataArray + """Create a test market dataset with random supply/demand values. - from muse.commodities import is_consumable, is_enduse, is_other + Args: + market: Template market dataset with dimensions + comm_usage: Sequence of commodity usage flags + Returns: + Dataset with supply, consumption and prices + """ shape = ( len(market.year), len(market.commodity), len(market.region), len(market.timeslice), ) + + values = np.random.randint(0, 5, shape) / np.random.randint(1, 5, shape) single = DataArray( - randint(0, 5, shape) / randint(1, 5, shape), + values, dims=("year", "commodity", "region", "timeslice"), coords={ "year": market.year, @@ -83,100 +87,102 @@ def sector_market(market: Dataset, comm_usage: Sequence[CommodityUsage]) -> Data def test_find_equilibrium(market: Dataset): - from copy import deepcopy - from unittest.mock import patch - - from numpy.random import choice - from pytest import approx - from xarray import broadcast - - from muse.commodities import is_enduse, is_other - from muse.mca import find_equilibrium + """Test market equilibrium finding with mock sectors. + Tests convergence behavior with different iteration limits and + verifies market values match expected outcomes. + """ market = market.interp(year=[2010, 2015]) - a_enduses = choice(market.commodity.values, 5, replace=False).tolist() + + # Setup test commodities + a_enduses = np.random.choice(market.commodity.values, 5, replace=False).tolist() b_enduses = [a_enduses.pop(), a_enduses.pop()] - # only "service" is currently truly meaningful and required here - # "service" means non-environmental outputs. + # Define commodity usage patterns available = ( CommodityUsage.CONSUMABLE, CommodityUsage.PRODUCT | CommodityUsage.ENVIRONMENTAL, CommodityUsage.OTHER, ) - a_usage = [ - CommodityUsage.PRODUCT if i in a_enduses else choice(available) - for i in market.commodity - ] - b_usage = [ - CommodityUsage.PRODUCT if i in b_enduses else choice(available) - for i in market.commodity - ] + def get_usage(enduses, commodities): + return [ + CommodityUsage.PRODUCT if c in enduses else np.random.choice(available) + for c in commodities + ] + + a_usage = get_usage(a_enduses, market.commodity) + b_usage = get_usage(b_enduses, market.commodity) + + # Create test markets a_market = sector_market(market, a_usage).rename(prices="costs") b_market = sector_market(market, b_usage).rename(prices="costs") + # Initialize market values market["supply"][:] = 0 market["consumption"][:] = 0 market["prices"][:] = 1 - cls = "muse.sectors.AbstractSector" - with patch(cls) as SectorA, patch(cls) as SectorB: - a = SectorA() - - side_effect_a = [0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0] - a.next.side_effect = lambda *args, **kwargs: a_market.sel( - commodity=~is_other(a_market.comm_usage) - ) * side_effect_a.pop(0) - - b = SectorB() - side_effect_b = [0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0] - b.next.side_effect = lambda *args, **kwargs: b_market.sel( - commodity=~is_other(b_market.comm_usage) - ) * side_effect_b.pop(0) - # maxiter equals 1 implicitly not convergence - result = find_equilibrium(market, deepcopy([a, b]), maxiter=2) - assert not result.converged - assert result.sectors[0].next.call_count == 1 - assert result.sectors[1].next.call_count == 1 - expected = a_market.supply + b_market.supply - actual, expected = broadcast(result.market.supply, expected) - assert actual.values == approx(0.7 * expected.values) - - side_effect_a.clear() - side_effect_a.extend([0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0]) - side_effect_b.clear() - side_effect_b.extend([0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0]) - result = find_equilibrium(market, deepcopy([a, b]), maxiter=5) - assert not result.converged - # check statelessness ~ only one call to next - assert result.sectors[0].next.call_count == 1 - assert result.sectors[1].next.call_count == 1 - expected = a_market.supply + b_market.supply - actual, expected = broadcast(result.market.supply, expected) - assert actual.values == approx(expected.values) - - side_effect_a.clear() - side_effect_a.extend([0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0]) - side_effect_b.clear() - side_effect_b.extend([0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0]) - sectors = deepcopy([a, b]) - result = find_equilibrium(market, sectors, maxiter=8) - assert result.converged - # check statelessness ~ only one call to next - assert result.sectors[0].next.call_count == 1 - assert result.sectors[1].next.call_count == 1 - - expected = a_market.supply + b_market.supply - actual, expected = broadcast(result.market.supply, expected) - assert actual.values == approx(expected.values) - - expected = a_market.consumption + b_market.consumption - actual, expected = broadcast(result.market.consumption, expected) - assert actual.values == approx(expected.values) - - expected = b_market.costs.where(is_enduse(b_market.comm_usage)) - expected = a_market.costs.where(is_enduse(a_market.comm_usage)) - expected = expected.where(expected > 1e-15, result.market.prices) - actual, expected = broadcast(result.market.prices, expected) - assert (actual.sel(year=2015)).values == approx(expected.sel(year=2015).values) + # Mock sector behavior + def create_mock_sector(test_market, side_effect): + sector = patch("muse.sectors.AbstractSector").start()() + sector.next.side_effect = lambda *args, **kwargs: ( + test_market.sel(commodity=~is_other(test_market.comm_usage)) + * side_effect.pop(0) + ) + return sector + + convergence_steps = [0.5, 0.7, 0.9, 0.95, 1.0, 1.0, 1.0] + + # Test with maxiter=2 (no convergence) + a = create_mock_sector(a_market, convergence_steps.copy()) + b = create_mock_sector(b_market, convergence_steps.copy()) + + result = find_equilibrium(market, deepcopy([a, b]), maxiter=2) + assert not result.converged + assert result.sectors[0].next.call_count == 1 + assert result.sectors[1].next.call_count == 1 + + expected = a_market.supply + b_market.supply + actual, expected = broadcast(result.market.supply, expected) + assert actual.values == approx(0.7 * expected.values) + + # Test with maxiter=5 (partial convergence) + a = create_mock_sector(a_market, convergence_steps.copy()) + b = create_mock_sector(b_market, convergence_steps.copy()) + + result = find_equilibrium(market, deepcopy([a, b]), maxiter=5) + assert not result.converged + assert all(s.next.call_count == 1 for s in result.sectors) + + actual, expected = broadcast( + result.market.supply, a_market.supply + b_market.supply + ) + assert actual.values == approx(expected.values) + + # Test with maxiter=8 (full convergence) + a = create_mock_sector(a_market, convergence_steps.copy()) + b = create_mock_sector(b_market, convergence_steps.copy()) + + result = find_equilibrium(market, deepcopy([a, b]), maxiter=8) + assert result.converged + assert all(s.next.call_count == 1 for s in result.sectors) + + # Verify final market state + actual, expected = broadcast( + result.market.supply, a_market.supply + b_market.supply + ) + assert actual.values == approx(expected.values) + + actual, expected = broadcast( + result.market.consumption, a_market.consumption + b_market.consumption + ) + assert actual.values == approx(expected.values) + + # Verify prices + expected = b_market.costs.where(is_enduse(b_market.comm_usage)) + expected = a_market.costs.where(is_enduse(a_market.comm_usage)) + expected = expected.where(expected > 1e-15, result.market.prices) + + actual, expected = broadcast(result.market.prices, expected) + assert actual.sel(year=2015).values == approx(expected.sel(year=2015).values) From 82c290f3c7b17782dd24e27a4bc8313daf660679 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:32:54 +0100 Subject: [PATCH 11/33] Tidy test_subsector --- tests/test_subsector.py | 110 ++++++++++++++++++++-------------------- 1 file changed, 55 insertions(+), 55 deletions(-) diff --git a/tests/test_subsector.py b/tests/test_subsector.py index efcbc1fb..358ffa8e 100644 --- a/tests/test_subsector.py +++ b/tests/test_subsector.py @@ -1,6 +1,16 @@ +from copy import deepcopy + import xarray as xr from pytest import fixture +from muse import constraints as cs +from muse import demand_share as ds +from muse import examples +from muse.agents.factories import create_agent +from muse.readers import read_csv_agent_parameters, read_initial_assets +from muse.readers.toml import read_settings +from muse.sectors.subsector import Subsector, aggregate_enduses + @fixture def model() -> str: @@ -9,24 +19,51 @@ def model() -> str: @fixture def technologies(model) -> xr.Dataset: - from muse import examples - return examples.sector("residential", model=model).technologies @fixture def market(model) -> xr.Dataset: - from muse import examples - return examples.residential_market(model) -def test_subsector_investing_aggregation(): - from copy import deepcopy +@fixture +def base_market(technologies, market): + """Common market setup used across tests.""" + return market.sel( + commodity=technologies.commodity, region=technologies.region + ).interp(year=[2020, 2025]) - from muse import examples - from muse.sectors.subsector import Subsector, aggregate_enduses +@fixture +def agent_params(model, tmp_path, technologies): + """Common agent parameters setup.""" + examples.copy_model(model, tmp_path) + path = tmp_path / "model" / "Agents.csv" + params = read_csv_agent_parameters(path) + capa = read_initial_assets(path.with_name("residential") / "ExistingCapacity.csv") + + for param in params: + param.update( + { + "capacity": deepcopy(capa.sel(region=param["region"])) + if param["agent_type"] == "retrofit" + else xr.zeros_like(capa.sel(region=param["region"])), + "agent_type": "default", + "category": "trade", + "year": 2020, + "search_rules": "from_assets -> compress -> reduce_assets", + "objectives": "ALCOE", + "decision": {"name": "mean", "parameters": ("ALCOE", False, 1)}, + } + ) + param.pop("quantity", None) + param.pop("share", None) + + return params + + +def test_subsector_investing_aggregation(): model_list = ["default", "medium"] sector_list = ["residential", "power", "gas"] @@ -40,53 +77,23 @@ def test_subsector_investing_aggregation(): market = mca.market.sel( commodity=technologies.commodity, region=technologies.region ).interp(year=[2020, 2025]) - subsector = Subsector(agents, commodities) + initial_agents = deepcopy(agents) + subsector = Subsector(agents, commodities) + assert {agent.year for agent in agents} == {int(market.year.min())} subsector.aggregate_lp(technologies.sel(year=2020), market) assert {agent.year for agent in agents} == {int(market.year.min() + 5)} + for initial, final in zip(initial_agents, agents): assert initial.assets.sum() != final.assets.sum() -def test_subsector_noninvesting_aggregation(market, model, technologies, tmp_path): - """Create some default agents and run subsector. - - Mostly a smoke test to check the returns look about right, with the right type and - containing "agent" dimensions. - """ - from copy import deepcopy - - from muse import constraints as cs - from muse import demand_share as ds - from muse import examples, readers - from muse.agents.factories import create_agent - from muse.sectors.subsector import Subsector, aggregate_enduses - - examples.copy_model(model, tmp_path) - path = tmp_path / "model" / "Agents.csv" - params = readers.read_csv_agent_parameters(path) - capa = readers.read_initial_assets( - path.with_name("residential") / "ExistingCapacity.csv" - ) - - for param in params: - if param["agent_type"] == "retrofit": - param["capacity"] = deepcopy(capa.sel(region=param["region"])) - else: - param["capacity"] = xr.zeros_like(capa.sel(region=param["region"])) - - if "share" in param: - del param["share"] - - param["agent_type"] = "default" - param["category"] = "trade" - param["year"] = 2020 - param["search_rules"] = "from_assets -> compress -> reduce_assets" - param["objectives"] = "ALCOE" - param["decision"]["parameters"] = ("ALCOE", False, 1) - param.pop("quantity") - agents = [create_agent(technologies=technologies, **param) for param in params] +def test_subsector_noninvesting_aggregation(base_market, agent_params, technologies): + """Test non-investing aggregation with default agents.""" + agents = [ + create_agent(technologies=technologies, **param) for param in agent_params + ] commodities = aggregate_enduses(technologies) subsector = Subsector( @@ -96,18 +103,11 @@ def test_subsector_noninvesting_aggregation(market, model, technologies, tmp_pat constraints=cs.factory("demand"), ) - market = market.sel( - commodity=technologies.commodity, region=technologies.region - ).interp(year=[2020, 2025]) assert all(agent.year == 2020 for agent in agents) - subsector.aggregate_lp(technologies.sel(year=2020), market) + subsector.aggregate_lp(technologies.sel(year=2020), base_market) def test_factory_smoke_test(model, technologies, tmp_path): - from muse import examples - from muse.readers.toml import read_settings - from muse.sectors.subsector import Subsector - examples.copy_model(model, tmp_path) settings = read_settings(tmp_path / "model" / "settings.toml") From f950683fec1cd4b24318e297b65df5c037b60f94 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:34:22 +0100 Subject: [PATCH 12/33] Tidy test_regressions --- tests/test_regressions.py | 84 ++++++++++++++++++--------------------- 1 file changed, 39 insertions(+), 45 deletions(-) diff --git a/tests/test_regressions.py b/tests/test_regressions.py index 028293a3..d62dfa00 100644 --- a/tests/test_regressions.py +++ b/tests/test_regressions.py @@ -1,59 +1,58 @@ +import numpy as np +import xarray as xr from pytest import approx, fixture +from muse.regressions import Exponential, Linear -@fixture -def regression_params(): - from numpy.random import rand - from xarray import Dataset - params = Dataset() - params["region"] = "region", ["USA", "ATE"] - params["sector"] = "sector", ["Residential", "Commercial"] - params["commodity"] = "commodity", ["algae", "agrires", "coal"] +def create_dataset(coords, random_vars): + """Helper to create test datasets with random variables.""" + ds = xr.Dataset(coords=coords) + shape = {k: len(v) for k, v in ds.coords.items()} + dims = tuple(shape.keys()) + + for var in random_vars: + ds[var] = dims, np.random.rand(*shape.values()) + return ds + - shape = {k: len(v) for k, v in params.coords.items()} - params["a"] = tuple(shape.keys()), rand(*shape.values()) - params["b"] = tuple(shape.keys()), rand(*shape.values()) - params["b0"] = tuple(shape.keys()), rand(*shape.values()) - params["b1"] = tuple(shape.keys()), rand(*shape.values()) - params["c"] = tuple(shape.keys()), rand(*shape.values()) - params["w"] = tuple(shape.keys()), rand(*shape.values()) - return params +@fixture +def regression_params(): + coords = { + "region": ["USA", "ATE"], + "sector": ["Residential", "Commercial"], + "commodity": ["algae", "agrires", "coal"], + } + random_vars = ["a", "b", "b0", "b1", "c", "w"] + return create_dataset(coords, random_vars) @fixture def drivers(): - from numpy.random import rand - from xarray import Dataset - - drivers = Dataset() - drivers["year"] = "year", [2010, 2015, 2020] - drivers["region"] = "region", ["USA", "ATE"] - drivers["gdp"] = ( - ("region", "year"), - (1000 * rand(len(drivers.region), len(drivers.year))), - ) - drivers["population"] = ( - ("region", "year"), - (10 * rand(len(drivers.region), len(drivers.year))), - ) - return drivers + coords = {"year": [2010, 2015, 2020], "region": ["USA", "ATE"]} + ds = xr.Dataset(coords=coords) + shape = (len(ds.region), len(ds.year)) + ds["gdp"] = ("region", "year"), 1000 * np.random.rand(*shape) + ds["population"] = ("region", "year"), 10 * np.random.rand(*shape) + return ds def test_exponential(regression_params, drivers): from numpy import exp from xarray import broadcast - from muse.regressions import Exponential - + # Prepare regression parameters and create functor rp = regression_params.drop_vars(("c", "w", "b0", "b1")) - functor = Exponential(**(rp.data_vars)) + functor = Exponential(**rp.data_vars) + + # Calculate expected and actual results actual = functor(drivers) factor = 1e6 * drivers.population * rp.a expected = factor * exp(drivers.population / drivers.gdp * rp.b) expected, actual = broadcast(expected, actual) assert actual.values == approx(expected.values) + # Test partial selection partial = actual.sel(region="USA") a, b = broadcast(partial, functor(drivers, region="USA")) assert a.values == approx(b.values) @@ -62,12 +61,12 @@ def test_exponential(regression_params, drivers): def test_linear(regression_params, drivers): from xarray import DataArray, broadcast - from muse.regressions import Linear - + # Prepare regression parameters and create functor rp = regression_params.drop_vars(("c", "w", "b")) functor = Linear(**rp.data_vars) - actual = functor(drivers, forecast=2) + # Test basic functionality + actual = functor(drivers, forecast=2) offset = drivers.gdp.sel(year=2010) / drivers.population.sel(year=2010) expected = rp.a * drivers.population + rp.b0 * ( drivers.gdp - offset * drivers.population @@ -75,6 +74,7 @@ def test_linear(regression_params, drivers): actual, expected = broadcast(actual, expected) assert actual.values == approx(expected.values) + # Test with interpolation year = [2010, 2012, 2014, 2020] scale = rp.b0.where( DataArray(year, coords={"year": year}, dims="year") + 2 < 2015, rp.b1 @@ -88,11 +88,5 @@ def test_linear(regression_params, drivers): if __name__ == "__main__": - from pytest import approx - - from muse import DEFAULT_SECTORS_DIRECTORY - from tests.agents import test_regressions - - sectors_dir = DEFAULT_SECTORS_DIRECTORY - regression_params = test_regressions.regression_params() - drivers = test_regressions.drivers() + regression_params = regression_params() + drivers = drivers() From dee4c7f9c50f56a596c9dceca4f638e93e3b5e6b Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:35:17 +0100 Subject: [PATCH 13/33] Tidy test_aggregoutput --- tests/test_aggregoutput.py | 140 ++++++++++++++----------------------- 1 file changed, 52 insertions(+), 88 deletions(-) diff --git a/tests/test_aggregoutput.py b/tests/test_aggregoutput.py index f11e8e74..d55a7450 100644 --- a/tests/test_aggregoutput.py +++ b/tests/test_aggregoutput.py @@ -1,120 +1,84 @@ -from muse import examples -from muse.outputs.mca import sector_capacity - - -def test_aggregate_sector(): - """Test for aggregate_sector function. +from pandas import DataFrame, concat - Check column titles, number of agents/region/technologies and assets capacities. - """ - from pandas import DataFrame, concat - - mca = examples.model("multiple_agents", test=True) - year = [2020, 2025] - sector_list = [sector for sector in mca.sectors if "preset" not in sector.name] - agent_list = [list(a.agents) for a in sector_list] - alldata = sector_capacity(sector_list[0]) +from muse import examples +from muse.outputs.mca import _aggregate_sectors, sector_capacity - columns = ["region", "agent", "type", "sector", "capacity"] +def _create_test_data(agents, years, sector_name): + """Helper function to create test DataFrame from agent data.""" frame = DataFrame() - for ai in agent_list[0]: - for y in year: - if y in ai.assets.year: - if ai.assets.capacity.sel(year=y).values > 0.0: + for agent in agents: + for year in years: + if year in agent.assets.year: + capacity = agent.assets.capacity.sel(year=year).values + if capacity > 0.0: data = DataFrame( { - "region": ai.region, - "agent": ai.name, - "type": ai.category, - "sector": sector_list[0].name, - "capacity": ai.assets.capacity.sel(year=y).values[0], + "region": agent.region, + "agent": agent.name, + "type": agent.category, + "sector": sector_name, + "capacity": capacity[0], }, - index=[(y, ai.assets.technology.values[0])], + index=[(year, agent.assets.technology.values[0])], ) frame = concat([frame, data]) + return frame - assert (frame[columns].values == alldata[columns].values).all() +def test_aggregate_sector(): + """Test aggregate_sector function with single sector data.""" + mca = examples.model("multiple_agents", test=True) + years = [2020, 2025] + sector_list = [sector for sector in mca.sectors if "preset" not in sector.name] + alldata = sector_capacity(sector_list[0]) -def test_aggregate_sectors(): - """Test for aggregate_sectors function.""" - from pandas import DataFrame, concat + frame = _create_test_data( + agents=list(sector_list[0].agents), years=years, sector_name=sector_list[0].name + ) - from muse.outputs.mca import _aggregate_sectors + columns = ["region", "agent", "type", "sector", "capacity"] + assert (frame[columns].values == alldata[columns].values).all() + +def test_aggregate_sectors(): + """Test aggregate_sectors function with multiple sectors.""" mca = examples.model("multiple_agents", test=True) - year = [2020, 2025, 2030] + years = [2020, 2025, 2030] sector_list = [sector for sector in mca.sectors if "preset" not in sector.name] - agent_list = [list(a.agents) for a in sector_list] alldata = _aggregate_sectors(mca.sectors, op=sector_capacity) - columns = ["region", "agent", "type", "sector", "capacity"] - frame = DataFrame() - for a, ai in enumerate(agent_list): - for ii in range(0, len(ai)): - for y in year: - if y in ai[ii].assets.year: - if ai[ii].assets.capacity.sel(year=y).values > 0.0: - data = DataFrame( - { - "region": ai[ii].region, - "agent": ai[ii].name, - "type": ai[ii].category, - "sector": sector_list[a].name, - "capacity": ai[ii] - .assets.capacity.sel(year=y) - .values[0], - }, - index=[(y, ai[ii].assets.technology.values[0])], - ) - frame = concat([frame, data]) + for sector in sector_list: + sector_frame = _create_test_data( + agents=list(sector.agents), years=years, sector_name=sector.name + ) + frame = concat([frame, sector_frame]) + columns = ["region", "agent", "type", "sector", "capacity"] assert (frame[columns].values == alldata[columns].values).all() def test_aggregate_sector_manyregions(): - """Test for aggregate_sector function with two regions. - - Check column titles, number of agents/region/technologies and assets capacities. - """ - from pandas import DataFrame, concat - - from muse.outputs.mca import _aggregate_sectors - + """Test aggregate_sector function with multiple regions.""" mca = examples.model("multiple_agents", test=True) residential = next(sector for sector in mca.sectors if sector.name == "residential") - agents = list(residential.agents) - agents[0].assets["region"] = "BELARUS" - agents[1].assets["region"] = "BELARUS" - agents[0].region = "BELARUS" - agents[1].region = "BELARUS" - year = [2020, 2025, 2030] + + # Set up Belarus region for testing + for agent in list(residential.agents)[:2]: + agent.assets["region"] = "BELARUS" + agent.region = "BELARUS" + + years = [2020, 2025, 2030] sector_list = [sector for sector in mca.sectors if "preset" not in sector.name] - agent_list = [list(a.agents) for a in sector_list] alldata = _aggregate_sectors(mca.sectors, op=sector_capacity) - columns = ["region", "agent", "type", "sector", "capacity"] - frame = DataFrame() - for a, ai in enumerate(agent_list): - for ii in range(0, len(ai)): - for y in year: - if y in ai[ii].assets.year: - if ai[ii].assets.capacity.sel(year=y).values > 0.0: - data = DataFrame( - { - "region": ai[ii].region, - "agent": ai[ii].name, - "type": ai[ii].category, - "sector": sector_list[a].name, - "capacity": ai[ii] - .assets.capacity.sel(year=y) - .values[0], - }, - index=[(y, ai[ii].assets.technology.values[0])], - ) - frame = concat([frame, data]) + for sector in sector_list: + sector_frame = _create_test_data( + agents=list(sector.agents), years=years, sector_name=sector.name + ) + frame = concat([frame, sector_frame]) + columns = ["region", "agent", "type", "sector", "capacity"] assert (frame[columns].values == alldata[columns].values).all() From fb5b3916cbd75bf166e2e48ec88198e9b7e96827 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:35:50 +0100 Subject: [PATCH 14/33] Tidy test_investments --- tests/test_investments.py | 39 +++++++++++++++++++-------------------- 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/tests/test_investments.py b/tests/test_investments.py index d22909a8..8543aa30 100644 --- a/tests/test_investments.py +++ b/tests/test_investments.py @@ -1,55 +1,54 @@ +import numpy as np +import xarray as xr from pytest import mark +from muse.investments import cliff_retirement_profile -def test_cliff_retirement_known_profile(): - from numpy import array - from xarray import DataArray - - from muse.investments import cliff_retirement_profile +def test_cliff_retirement_known_profile(): + """Test cliff retirement profile with known values and expected output.""" technology = ["a", "b", "c"] - lifetime = DataArray( - range(1, 1 + len(technology)), + lifetime = xr.DataArray( + np.arange(1, len(technology) + 1), dims="technology", coords={"technology": technology}, name="technical_life", ) profile = cliff_retirement_profile(technical_life=lifetime, investment_year=2020) - expected = array( + expected = np.array( [ [True, False, False, False], [True, True, False, False], [True, True, True, False], ] ) + assert set(profile.dims) == {"year", "technology"} assert (profile == expected.T).all() @mark.parametrize("protected", range(12)) def test_cliff_retirement_random_profile(protected): - from numpy.random import randint - from xarray import DataArray - - from muse.investments import cliff_retirement_profile - + """Test cliff retirement profile with random lifetimes and protected periods.""" technology = list("abcde") - - lifetime = DataArray( - sorted(randint(1, 10, len(technology))), + lifetime = xr.DataArray( + sorted(np.random.randint(1, 10, len(technology))), dims="technology", coords={"technology": technology}, name="technical_life", ) - effective_lifetime = (protected // lifetime + 1) * lifetime investment_year = 2020 + effective_lifetime = (protected // lifetime + 1) * lifetime profile = cliff_retirement_profile( technical_life=lifetime.clip(min=protected), investment_year=investment_year ) + + # Verify profile boundaries and properties + profile_int = profile.astype(int) assert profile.year.min() == investment_year assert profile.year.max() <= investment_year + effective_lifetime.max() + 1 - assert profile.astype(int).interp(year=investment_year).all() - assert profile.astype(int).interp(year=investment_year + protected - 1).all() - assert not profile.astype(int).interp(year=profile.year.max()).any() + assert profile_int.interp(year=investment_year).all() + assert profile_int.interp(year=investment_year + protected - 1).all() + assert not profile_int.interp(year=profile.year.max()).any() From 7b137ba5c661dfc21d1aa9f13877640e3ba1eeba Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:38:05 +0100 Subject: [PATCH 15/33] Tidy test_interactions --- tests/test_interactions.py | 108 ++++++++++++++++++++----------------- 1 file changed, 58 insertions(+), 50 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index e2d4a490..b91c2ed9 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -1,82 +1,90 @@ """Test agent interactions.""" +from collections import namedtuple +from itertools import chain + import pytest +from numpy.random import choice from pytest import fixture, mark +from muse.errors import NoInteractionsFound +from muse.interactions import ( + agents_groupby, + factory, + new_to_retro_net, + register_agent_interaction, +) + @fixture def agents(): - from collections import namedtuple - - from numpy.random import choice - Agent = namedtuple("Agent", ["region", "name", "category", "assets"]) - regions = ["Area52", "Bermuda Triangle", "City of London"] - names = ["John", "Joe", "Jill"] - categories = ["yup", "nope"] - results = { - (region, name, cat) - for region, name, cat in zip( - *(choice(data, 40) for data in [regions, names, categories]) - ) + sample_data = { + "regions": ["Area52", "Bermuda Triangle", "City of London"], + "names": ["John", "Joe", "Jill"], + "categories": ["yup", "nope"], } - return [Agent(*u, [None]) for u in results] + # Generate unique combinations of region, name, and category + combinations = { + tuple(items) + for items in zip(*(choice(data, 40) for data in sample_data.values())) + } -def test_groupby(agents): - from itertools import chain + return [Agent(*combo, [None]) for combo in combinations] - from muse.interactions import agents_groupby +def test_groupby(agents): grouped = agents_groupby(agents, ("category", "name")) - assert sum(len(u) for u in grouped.values()) == len(agents) - assert set(id(u) for u in chain(*grouped.values())) == set(id(u) for u in agents) - assert set(grouped.keys()) == set((u.category, u.name) for u in agents) - for (category, name), group_agents in grouped.items(): - assert all(agent.category == category for agent in group_agents) - assert all(agent.name == name for agent in group_agents) + # Verify group sizes and agent preservation + assert sum(len(group) for group in grouped.values()) == len(agents) + assert set(map(id, chain(*grouped.values()))) == set(map(id, agents)) -def test_new_to_retro_net(agents): - from itertools import chain + # Verify correct grouping keys and contents + assert set(grouped.keys()) == {(a.category, a.name) for a in agents} + for (cat, name), group in grouped.items(): + assert all(a.category == cat and a.name == name for a in group) - from muse.interactions import new_to_retro_net +def test_new_to_retro_net(agents): net = new_to_retro_net(agents, "nope") - assert sum(len(u) for u in net) <= len(agents) - assert set(id(u) for u in chain(*net)).issubset(id(u) for u in agents) - assert all(len(list(u)) == 2 for u in net) - for agents in net: - assert len(set(u.region for u in agents)) == 1 - assert len(set(u.name for u in agents)) == 1 - categories = [u.category for u in agents] - i = categories.index("yup") if "yup" in categories else 0 - assert "nope" not in categories[i:] - assert "yup" not in categories[:i] + + # Verify network properties + assert sum(len(group) for group in net) <= len(agents) + assert set(map(id, chain(*net))).issubset(map(id, agents)) + assert all(len(list(group)) == 2 for group in net) + + # Verify group constraints + for group in net: + assert len({a.region for a in group}) == 1 + assert len({a.name for a in group}) == 1 + categories = [a.category for a in group] + yup_index = categories.index("yup") if "yup" in categories else 0 + assert "nope" not in categories[yup_index:] + assert "yup" not in categories[:yup_index] @mark.usefixtures("save_registries") def test_compute_interactions(agents): - from muse.errors import NoInteractionsFound - from muse.interactions import factory, new_to_retro_net, register_agent_interaction - @register_agent_interaction def dummy_interaction(a, b): - assert a.assets[0] is None - assert b.assets[0] is None - a.assets[0] = b - b.assets[0] = a + assert all(agent.assets[0] is None for agent in (a, b)) + a.assets[0], b.assets[0] = b, a interactions = factory([("new_to_retro", "dummy_interaction")]) interactions(agents) - are_none = [agent for agent in agents if agent.assets[0] is None] - assert len({(u.region, u.name) for u in are_none}) == len(are_none) - not_none = [agent for agent in agents if agent.assets[0] is not None] - assert len(new_to_retro_net(agents)) == 0 or len(not_none) != 0 - for agent in not_none: - assert agent.assets[0].assets[0] is agent + # Check unmatched agents + unmatched = [agent for agent in agents if agent.assets[0] is None] + assert len({(a.region, a.name) for a in unmatched}) == len(unmatched) + + # Check matched agents + matched = [agent for agent in agents if agent.assets[0] is not None] + assert len(new_to_retro_net(agents)) == 0 or matched + assert all(agent.assets[0].assets[0] is agent for agent in matched) - agents2 = [a for a in agents if a.category == "nope"] + # Test that NoInteractionsFound is raised when all agents are 'nope' category + nope_agents = [a for a in agents if a.category == "nope"] with pytest.raises(NoInteractionsFound): - interactions(agents2) + interactions(nope_agents) From f81993c5ea4b3fa5fbec7de164507422a78a0fdb Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:40:27 +0100 Subject: [PATCH 16/33] Tidy test_commodity --- tests/test_commodity.py | 40 ++++++++++++++++++++++++++-------------- 1 file changed, 26 insertions(+), 14 deletions(-) diff --git a/tests/test_commodity.py b/tests/test_commodity.py index 1c1aadff..c990d5f4 100644 --- a/tests/test_commodity.py +++ b/tests/test_commodity.py @@ -1,27 +1,39 @@ def test_from_technologies(technologies, coords): + """Test CommodityUsage.from_technologies method. + + Verifies that commodity usage flags are correctly set based on technology + inputs/outputs and commodity types. + """ from muse.commodities import CommodityUsage technologies = technologies.drop_vars("comm_usage") technologies["comm_type"] = "commodity", coords["comm_type"] technologies = technologies.set_coords("comm_type") + comm_usage = CommodityUsage.from_technologies(technologies) - redux = ( + # Check which commodities are used in any region/year/technology + usage_mask = ( technologies[["fixed_inputs", "fixed_outputs", "flexible_inputs"]] > 0 ).any(("region", "year", "technology")) - for actual, is_cons, is_prod in zip( - comm_usage, redux.fixed_inputs | redux.flexible_inputs, redux.fixed_outputs + # Test individual commodity usage flags + for actual, is_consumable, is_product in zip( + comm_usage, + usage_mask.fixed_inputs | usage_mask.flexible_inputs, + usage_mask.fixed_outputs, ): - assert bool(actual & CommodityUsage.PRODUCT) == is_prod - assert bool(actual & CommodityUsage.CONSUMABLE) == is_cons - - assert ((comm_usage & CommodityUsage.PRODUCT != 0) == redux.fixed_outputs).all() - assert ( - (comm_usage & CommodityUsage.ENVIRONMENTAL != 0) - == (redux.comm_type == "environmental") - ).all() + assert bool(actual & CommodityUsage.PRODUCT) == is_product + assert bool(actual & CommodityUsage.CONSUMABLE) == is_consumable - assert ( - (comm_usage & CommodityUsage.ENERGY != 0) == (redux.comm_type == "energy") - ).all() + # Test commodity type flags across all items + for i in range(len(comm_usage)): + assert ( + bool(comm_usage[i] & CommodityUsage.PRODUCT) == usage_mask.fixed_outputs[i] + ) + assert bool(comm_usage[i] & CommodityUsage.ENVIRONMENTAL) == ( + usage_mask.comm_type[i] == "environmental" + ) + assert bool(comm_usage[i] & CommodityUsage.ENERGY) == ( + usage_mask.comm_type[i] == "energy" + ) From 1a3a676a99d1a0446154e650e3177944716fc316 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:42:13 +0100 Subject: [PATCH 17/33] Tidy test_utilities --- tests/test_utilities.py | 261 ++++++++++++++++++++-------------------- 1 file changed, 132 insertions(+), 129 deletions(-) diff --git a/tests/test_utilities.py b/tests/test_utilities.py index 53eababc..93eed317 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -4,25 +4,39 @@ def make_array(array): + """Create a random DataArray with the same dimensions and coordinates as input.""" data = np.random.randint(1, 5, len(array)) * ( np.random.randint(0, 10, len(array)) > 8 ) return xr.DataArray(data, dims=array.dims, coords=array.coords) +def assert_shape_and_dims(actual, expected_shape, expected_dims=None): + """Helper to verify array shape and dimensions. + + Works with both numpy arrays and xarray objects. + For numpy arrays, only shape is checked. + For xarray objects, both shape and dimensions are checked. + """ + assert actual.ndim == len(expected_shape) + assert actual.shape == expected_shape + if expected_dims is not None and hasattr(actual, "dims"): + assert set(actual.dims) == set(expected_dims) + + @mark.parametrize( "coordinates", [("technology", "installed", "region"), ("technology", "installed"), ("region",)], ) def test_reduce_assets(coordinates: tuple, capacity: xr.DataArray): + """Test reducing assets along specified coordinates.""" from muse.utilities import reduce_assets actual = reduce_assets(capacity, coords=coordinates) uniques = set(zip(*(getattr(capacity, d).values for d in coordinates))) assert len(uniques) == len(actual.asset) - actual_uniques = set(zip(*(getattr(actual, d).values for d in coordinates))) - assert uniques == actual_uniques + assert uniques == set(zip(*(getattr(actual, d).values for d in coordinates))) for index in actual.asset: condition = True @@ -33,6 +47,7 @@ def test_reduce_assets(coordinates: tuple, capacity: xr.DataArray): def test_reduce_assets_with_zero_size(capacity: xr.DataArray): + """Test reducing assets with empty selection.""" from muse.utilities import reduce_assets x = capacity.sel(asset=[]) @@ -41,14 +56,13 @@ def test_reduce_assets_with_zero_size(capacity: xr.DataArray): def test_broadcast_over_assets(technologies, capacity): + """Test broadcasting over assets with different year settings.""" from muse.utilities import broadcast_over_assets - # Test with installed_as_year = True result1 = broadcast_over_assets(technologies, capacity, installed_as_year=True) assert set(result1.dims) == {"asset", "commodity"} assert (result1.asset == capacity.asset).all() - # Test with installed_as_year = False result2 = broadcast_over_assets(technologies, capacity, installed_as_year=False) assert set(result2.dims) == {"asset", "commodity", "year"} assert (result2.asset == capacity.asset).all() @@ -58,150 +72,138 @@ def test_broadcast_over_assets(technologies, capacity): # broadcast_over_assets(technologies, technologies) +def create_test_array(shape, scale=20, offset=-10): + """Create a random integer array for testing.""" + return (np.random.rand(*shape) * scale + offset).astype(int) + + def test_tupled_dimension_no_tupling(): + """Test tupled dimension conversion without actual tupling.""" from muse.utilities import tupled_dimension - array = (np.random.rand(10, 1) * 20 - 10).astype(int) + array = create_test_array((10, 1)) actual = tupled_dimension(array, 1) - assert actual.ndim == 1 - assert actual.shape == array.shape[:-1] + assert_shape_and_dims(actual, array.shape[:-1]) assert actual == approx(array.reshape(array.shape[0])) - array = (np.random.rand(10, 1, 5) * 20 - 10).astype(int) + array = create_test_array((10, 1, 5)) actual = tupled_dimension(array, 1) - assert actual.ndim == 2 - assert actual.shape == (array.shape[0], array.shape[2]) + assert_shape_and_dims(actual, (array.shape[0], array.shape[2])) assert actual == approx(array.reshape(array.shape[0], array.shape[2])) +def verify_tupled_output(actual, array, axis, indices): + """Helper to verify tupled dimension output.""" + for idx in indices: + if len(idx) == 1: + i = idx[0] + assert len(actual[i]) == array.shape[axis] + assert isinstance(actual[i], tuple) + assert tuple(array[i, :] if axis == 1 else array[:, i]) == actual[i] + else: + i, j = idx + assert len(actual[i, j]) == array.shape[axis] + assert isinstance(actual[i, j], tuple) + if axis == 0: + assert tuple(array[:, i, j]) == actual[i, j] + elif axis == 1: + assert tuple(array[i, :, j]) == actual[i, j] + else: + assert tuple(array[i, j, :]) == actual[i, j] + + def test_tupled_dimension_2d(): + """Test tupled dimension conversion for 2D arrays.""" from muse.utilities import tupled_dimension - array = (np.random.rand(10, 3) * 20 - 10).astype(int) + array = create_test_array((10, 3)) - actual = tupled_dimension(array, 1) - assert actual.ndim == 1 - assert actual.shape == (array.shape[0],) - for i in range(array.shape[0]): - assert len(actual[i]) == array.shape[1] - assert isinstance(actual[i], tuple) - assert tuple(array[i, :]) == actual[i] - - actual = tupled_dimension(array, 0) - assert actual.ndim == 1 - assert actual.shape == (array.shape[1],) - for i in range(array.shape[1]): - assert len(actual[i]) == array.shape[0] - assert isinstance(actual[i], tuple) - assert tuple(array[:, i]) == actual[i] + for axis in (0, 1): + actual = tupled_dimension(array, axis) + expected_shape = (array.shape[1],) if axis == 0 else (array.shape[0],) + assert_shape_and_dims(actual, expected_shape) + verify_tupled_output( + actual, array, axis, [(i,) for i in range(expected_shape[0])] + ) def test_tupled_dimension_3d(): + """Test tupled dimension conversion for 3D arrays.""" from muse.utilities import tupled_dimension - array = (np.random.rand(10, 3, 5) * 20 - 10).astype(int) - - actual = tupled_dimension(array, 1) - assert actual.ndim == 2 - assert actual.shape == (array.shape[0], array.shape[2]) - for i in range(array.shape[0]): - for j in range(array.shape[2]): - assert len(actual[i, j]) == array.shape[1] - assert isinstance(actual[i, j], tuple) - assert tuple(array[i, :, j]) == actual[i, j] - - actual = tupled_dimension(array, 0) - assert actual.ndim == 2 - assert actual.shape == (array.shape[1], array.shape[2]) - for i in range(array.shape[1]): - for j in range(array.shape[2]): - assert len(actual[i, j]) == array.shape[0] - assert isinstance(actual[i, j], tuple) - assert tuple(array[:, i, j]) == actual[i, j] - - actual = tupled_dimension(array, 2) - assert actual.ndim == 2 - assert actual.shape == (array.shape[0], array.shape[1]) - for i in range(array.shape[0]): - for j in range(array.shape[1]): - assert len(actual[i, j]) == array.shape[2] - assert isinstance(actual[i, j], tuple) - assert tuple(array[i, j, :]) == actual[i, j] - - -@mark.parametrize("order", [["a", "b", "c"], ["b", "c", "a"], ["c", "a", "b"]]) -def test_lexical_with_bin(order): - """Test lexical comparison against hand-constructed tuples.""" - from muse.utilities import lexical_comparison - + array = create_test_array((10, 3, 5)) + + for axis in (0, 1, 2): + actual = tupled_dimension(array, axis) + if axis == 0: + expected_shape = (array.shape[1], array.shape[2]) + elif axis == 1: + expected_shape = (array.shape[0], array.shape[2]) + else: + expected_shape = (array.shape[0], array.shape[1]) + assert_shape_and_dims(actual, expected_shape) + indices = [ + (i, j) for i in range(expected_shape[0]) for j in range(expected_shape[1]) + ] + verify_tupled_output(actual, array, axis, indices) + + +def create_test_objectives(shape=(5, 10)): + """Create test objectives dataset.""" objectives = xr.Dataset() - objectives["a"] = ("asset", "replacement"), np.random.rand(5, 10) * 10 - 5 - objectives["b"] = ("asset", "replacement"), np.random.rand(5, 10) * 10 - 5 - objectives["c"] = ("asset", "replacement"), np.random.rand(5, 10) * 10 - 5 + for var in ["a", "b", "c"]: + objectives[var] = ("asset", "replacement"), np.random.rand(*shape) * 10 - 5 objectives["asset"] = np.random.choice( objectives.replacement, len(objectives.asset), replace=False ) + return objectives - binsizes = xr.Dataset( + +def create_test_binsizes(): + """Create test binsizes dataset.""" + return xr.Dataset( { "a": np.random.rand() * 0.1, "b": -np.random.rand() * 0.1, "c": np.random.rand(), } ) - expected = np.zeros(shape=objectives.a.shape, dtype=object) - for i in range(expected.shape[0]): - for j in range(expected.shape[1]): - expected[i, j] = ( - int(np.floor(objectives[order[0]][i, j] / binsizes[order[0]])), - int(np.floor(objectives[order[1]][i, j] / binsizes[order[1]])), - int(np.floor(objectives[order[2]][i, j] / binsizes[order[2]])), - ) - - actual = lexical_comparison(objectives, binsizes[order]) - assert actual.shape == expected.shape - for i in range(expected.shape[0]): - for j in range(expected.shape[1]): - assert actual.values[i, j] == expected[i, j] @mark.parametrize("order", [["a", "b", "c"], ["b", "c", "a"], ["c", "a", "b"]]) -def test_lexical_nobin(order): - """Test lexical comparison against hand-constructed tuples.""" +def test_lexical_comparison(order): + """Test lexical comparison with and without binning.""" from muse.utilities import lexical_comparison - objectives = xr.Dataset() - objectives["a"] = ("asset", "replacement"), np.random.rand(5, 10) * 10 - 5 - objectives["b"] = ("asset", "replacement"), np.random.rand(5, 10) * 10 - 5 - objectives["c"] = ("asset", "replacement"), np.random.rand(5, 10) * 10 - 5 - objectives["asset"] = np.random.choice( - objectives.replacement, len(objectives.asset), replace=False - ) - - binsizes = xr.Dataset( - { - "a": np.random.rand() * 0.1, - "b": -np.random.rand() * 0.1, - "c": np.random.rand(), - } - ) - expected = np.zeros(shape=objectives.a.shape, dtype=object) - for i in range(expected.shape[0]): - for j in range(expected.shape[1]): - expected[i, j] = ( - int(np.floor(objectives[order[0]][i, j] / binsizes[order[0]])), - int(np.floor(objectives[order[1]][i, j] / binsizes[order[1]])), - objectives[order[2]][i, j] / binsizes[order[2]], - ) + objectives = create_test_objectives() + binsizes = create_test_binsizes() + + def create_expected(bin_last=True): + expected = np.zeros(shape=objectives.a.shape, dtype=object) + for i in range(expected.shape[0]): + for j in range(expected.shape[1]): + values = [] + for k in range(3): + val = objectives[order[k]][i, j] / binsizes[order[k]] + values.append(int(np.floor(val)) if bin_last or k < 2 else val) + expected[i, j] = tuple(values) + return expected + + # Test with binning + actual = lexical_comparison(objectives, binsizes[order]) + expected = create_expected(bin_last=True) + assert actual.shape == expected.shape + assert (actual.values == expected).all() + # Test without binning last value actual = lexical_comparison(objectives, binsizes[order], bin_last=False) + expected = create_expected(bin_last=False) assert actual.shape == expected.shape - for i in range(expected.shape[0]): - for j in range(expected.shape[1]): - assert actual.values[i, j] == expected[i, j] + assert (actual.values == expected).all() def test_merge_assets(): + """Test merging assets with different coordinate orders.""" from numpy import arange from muse.utilities import interpolate_capacity, merge_assets @@ -215,31 +217,31 @@ def fake(year, order=("installed", "technology")): ("year", "asset"), np.random.rand(len(result.year), len(result.asset)), ) - result = result[["capacity", *order]].set_coords(order) - return result.capacity + return result[["capacity", *order]].set_coords(order).capacity - # checks order of coords does not interfere with merging order = ["installed", "technology"] capa_a = fake(np.arange(2010, 2020, 3, dtype="int64"), order) np.random.shuffle(order) capa_b = fake(arange(2014, 2024, 2, dtype="int64"), order) actual = merge_assets(capa_a, capa_b) - assert actual.installed.dtype == capa_a.installed.dtype - assert capa_a.installed.isin(actual.installed).all() - assert capa_b.installed.isin(actual.installed).all() - assert capa_a.technology.isin(actual.technology).all() - assert capa_b.technology.isin(actual.technology).all() - assert capa_a.year.isin(actual.year).all() - assert capa_b.year.isin(actual.year).all() + # Verify coordinate preservation + for coord in ["installed", "technology", "year"]: + assert getattr(capa_a, coord).isin(getattr(actual, coord)).all() + assert getattr(capa_b, coord).isin(getattr(actual, coord)).all() + # Verify asset uniqueness assets = [(i, t) for i, t in zip(actual.installed.values, actual.technology.values)] assert len(actual.asset) == len(set(assets)) - assets = [ + all_assets = [ (i, t) for i, t in zip(capa_a.installed.values, capa_a.technology.values) - ] + [(i, t) for i, t in zip(capa_b.installed.values, capa_b.technology.values)] - assert len(actual.asset) == len(set(assets)) + ] + all_assets.extend( + (i, t) for i, t in zip(capa_b.installed.values, capa_b.technology.values) + ) + assert len(actual.asset) == len(set(all_assets)) + # Verify capacity values for inst, tech in zip(actual.installed.values, actual.technology.values): ab_side = actual.sel( asset=((actual.installed == inst) & (actual.technology == tech)) @@ -260,6 +262,7 @@ def fake(year, order=("installed", "technology")): def test_avoid_repetitions(): + """Test avoiding repetitions in time series data.""" from muse.utilities import avoid_repetitions start, end = 2010, 2010 + 3 * 5 @@ -272,12 +275,11 @@ def test_avoid_repetitions(): np.random.randint(0, 10, (len(assets.year), len(assets.asset))), ) - assets.capacity.loc[{"year": list(range(start + 1, end, 3))}] = assets.capacity.sel( - year=list(range(start, end, 3)) - ).values - assets.capacity.loc[{"year": list(range(start + 2, end, 3))}] = assets.capacity.sel( - year=list(range(start, end, 3)) - ).values + # Create repetitions in the data + for offset in (1, 2): + assets.capacity.loc[{"year": list(range(start + offset, end, 3))}] = ( + assets.capacity.sel(year=list(range(start, end, 3))).values + ) result = assets.sel(year=avoid_repetitions(assets.capacity)) assert 3 * len(result.year) == 2 * len(assets.year) @@ -286,6 +288,7 @@ def test_avoid_repetitions(): def test_check_dimensions(): + """Test dimension checking functionality.""" from muse.utilities import check_dimensions data = xr.DataArray( @@ -294,13 +297,13 @@ def test_check_dimensions(): coords={"dim1": range(4), "dim2": range(5)}, ) - # Valid + # Test valid case check_dimensions(data, required=["dim1"], optional=["dim2"]) - # Missing required + # Test missing required dimension with raises(ValueError, match="Missing required dimensions"): check_dimensions(data, required=["dim1", "dim3"], optional=["dim2"]) - # Extra dimension + # Test extra dimension with raises(ValueError, match="Extra dimensions"): check_dimensions(data, required=["dim1"]) From 7dc6e5b4f3b8632ebe07c7ad36b91cd08de00903 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:44:24 +0100 Subject: [PATCH 18/33] Tidy test_filters --- tests/test_filters.py | 136 ++++++++++++++++++++---------------------- 1 file changed, 66 insertions(+), 70 deletions(-) diff --git a/tests/test_filters.py b/tests/test_filters.py index c98744a3..5a2899cb 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -1,8 +1,33 @@ +from collections import namedtuple + import numpy as np import xarray as xr from pytest import fixture, mark -from muse.filters import factory, register_filter, register_initializer +from muse.commodities import is_enduse +from muse.filters import ( + currently_existing_tech, + factory, + initialize_from_assets, + initialize_from_technologies, + maturity, + register_filter, + register_initializer, + same_enduse, + same_fuels, + similar_technology, +) + + +# Common test utilities +def assert_tech_comparison(actual, search_space, tech_attr, technologies): + """Helper for technology comparison tests.""" + assert sorted(actual.dims) == sorted(search_space.dims) + attr_values = getattr(technologies, tech_attr) + for tech in actual.replacement: + for asset in actual.asset: + expected = attr_values.loc[tech] == attr_values.loc[asset] + assert expected == actual.sel(replacement=tech, asset=asset) @fixture @@ -20,7 +45,6 @@ def search_space(retro_agent, technologies): @fixture def technologies(technologies): - # Filters must take technology data for a single year return technologies.sel(year=2010) @@ -51,9 +75,7 @@ def start(retro_agent, demand, **kwargs): @register_filter def first(retro_agent, search_space, switch=True, data=None): - if switch: - return search_space[2:] - return search_space[:2] + return search_space[2:] if switch else search_space[:2] @register_filter def second(retro_agent, search_space, switch=True, data=None): @@ -62,7 +84,6 @@ def second(retro_agent, search_space, switch=True, data=None): sp = start(None, None) assert factory(["start", "first"])(None, sp) == sp[2:] assert factory(["start", "first"])(None, sp, switch=False) == sp[:2] - assert factory(["start", "second"])(None, sp, data=(1, 3, 5)) == [1, 3] assert factory(["start", "first", "second"])(None, sp, data=(1, 3, 5)) == [3] assert factory(["start", "first", "second"])( @@ -71,9 +92,6 @@ def second(retro_agent, search_space, switch=True, data=None): def test_same_enduse(retro_agent, technologies, search_space): - from muse.commodities import is_enduse - from muse.filters import same_enduse - result = same_enduse(retro_agent, search_space, technologies=technologies) enduses = is_enduse(technologies.comm_usage) finputs = technologies.sel(region=retro_agent.region, commodity=enduses) @@ -81,11 +99,17 @@ def test_same_enduse(retro_agent, technologies, search_space): expected = search_space.copy() for asset in result.asset: - asset_enduses = finputs.sel(technology=asset) - asset_enduses = set(asset_enduses.commodity.loc[asset_enduses].values) + asset_enduses = set( + finputs.sel(technology=asset) + .commodity.loc[finputs.sel(technology=asset)] + .values + ) for tech in result.replacement: - tech_enduses = finputs.sel(technology=tech) - tech_enduses = set(tech_enduses.commodity.loc[tech_enduses].values) + tech_enduses = set( + finputs.sel(technology=tech) + .commodity.loc[finputs.sel(technology=tech)] + .values + ) expected.loc[{"replacement": tech, "asset": asset}] = ( asset_enduses.issubset(tech_enduses) ) @@ -95,34 +119,17 @@ def test_same_enduse(retro_agent, technologies, search_space): def test_similar_tech(retro_agent, search_space, technologies): - from muse.filters import similar_technology - actual = similar_technology(retro_agent, search_space, technologies=technologies) - assert sorted(actual.dims) == sorted(search_space.dims) - - tech_type = technologies.tech_type - for tech in actual.replacement: - for asset in actual.asset: - expected = tech_type.loc[tech] == tech_type.loc[asset] - assert expected == actual.sel(replacement=tech, asset=asset) + assert_tech_comparison(actual, search_space, "tech_type", technologies) def test_similar_fuels(retro_agent, search_space, technologies): - from muse.filters import same_fuels - actual = same_fuels(retro_agent, search_space, technologies=technologies) - assert sorted(actual.dims) == sorted(search_space.dims) - - fuel_type = technologies.fuel - for tech in actual.replacement: - for asset in actual.asset: - expected = fuel_type.loc[tech] == fuel_type.loc[asset] - assert expected == actual.sel(replacement=tech, asset=asset) + assert_tech_comparison(actual, search_space, "fuel", technologies) def test_currently_existing(retro_agent, search_space, technologies, agent_market, rng): - from muse.filters import currently_existing_tech - + # Test with zero capacity agent_market.capacity[:] = 0 actual = currently_existing_tech( retro_agent, search_space, technologies=technologies, market=agent_market @@ -130,15 +137,16 @@ def test_currently_existing(retro_agent, search_space, technologies, agent_marke assert sorted(actual.dims) == sorted(search_space.dims) assert not actual.any() + # Test with full capacity agent_market.capacity[:] = 1 actual = currently_existing_tech( retro_agent, search_space, technologies=technologies, market=agent_market ) - assert sorted(actual.dims) == sorted(search_space.dims) in_market = search_space.replacement.isin(agent_market.technology) assert not actual.sel(replacement=~in_market).any() assert actual.sel(replacement=in_market).all() + # Test with partial capacity techs = rng.choice( list(set(agent_market.technology.values)), 1 + rng.choice(range(len(set(agent_market.technology.values)))), @@ -149,7 +157,7 @@ def test_currently_existing(retro_agent, search_space, technologies, agent_marke actual = currently_existing_tech( retro_agent, search_space, technologies=technologies, market=agent_market ) - assert sorted(actual.dims) == sorted(search_space.dims) + assert not actual.sel(replacement=~in_market).any() current_cap = agent_market.capacity.sel( year=retro_agent.year, region=retro_agent.region @@ -160,58 +168,50 @@ def test_currently_existing(retro_agent, search_space, technologies, agent_marke @mark.xfail def test_maturity(retro_agent, search_space, technologies, agent_market): - from muse.commodities import is_enduse - from muse.filters import maturity - enduses = is_enduse(technologies.comm_usage) outputs = technologies.fixed_outputs.sel(commodity=enduses, region="USA", year=2010) capacity = agent_market.capacity.sel(year=2010, region="USA") production = (outputs * capacity).sum("technology") - # nothing should be true - retro_agent.maturity_threshold = 1.1 * (capacity / production).max() - actual = maturity(retro_agent, search_space, technologies, agent_market) - assert sorted(actual.dims) == sorted(search_space.dims) - assert (not actual).all() - - # some should be true - do it with a fully on search space for simplicity - retro_agent.maturity_threshold = 0.8 * (capacity / production).max() - actual = maturity( - search_space == retro_agent, search_space, technologies, agent_market - ) - assert sorted(actual.dims) == sorted(search_space.dims) - assert actual.any() - # all should be true - retro_agent.maturity_threshold = 0.8 * (capacity / production).min() - actual = maturity(retro_agent, search_space, technologies, agent_market) - assert (actual == search_space).any() + def check_maturity(threshold_factor, expected_result=None): + retro_agent.maturity_threshold = ( + threshold_factor * (capacity / production).max() + ) + actual = maturity(retro_agent, search_space, technologies, agent_market) + assert sorted(actual.dims) == sorted(search_space.dims) + if expected_result is not None: + assert (actual == expected_result).all() + return actual + + # Test different threshold scenarios + assert not check_maturity(1.1).any() # Nothing should be true + assert check_maturity(0.8).any() # Some should be true + assert ( + check_maturity( + 0.8 * (capacity / production).min() / (capacity / production).max() + ) + == search_space + ).any() def test_init_from_tech(demand_share, technologies, agent_market): - from collections import namedtuple - - from muse.filters import initialize_from_technologies - agent = namedtuple("DummyAgent", ["tolerance"])(tolerance=1e-8) - # All technologies produce demanded commodities + # Test with producing technologies space = initialize_from_technologies(agent, demand_share, technologies=technologies) assert set(space.dims) == {"asset", "replacement"} assert (space.asset.values == demand_share.asset.values).all() assert (space.replacement.values == technologies.technology.values).all() assert space.all() - # No technology produces demanded commodities + # Test with non-producing technologies technologies.fixed_outputs[:] = 0 space = initialize_from_technologies(agent, demand_share, technologies=technologies) assert not space.any() def test_init_from_asset(technologies, rng): - from collections import namedtuple - - from muse.filters import initialize_from_assets - + # Create test data technology = rng.choice(technologies.technology, 5) installed = rng.choice((2020, 2025), len(technology)) year = np.arange(2020, 2040, 5) @@ -227,6 +227,7 @@ def test_init_from_asset(technologies, rng): ) agent = namedtuple("DummyAgent", ["assets"])(xr.Dataset(dict(capacity=capacity))) + # Test with assets space = initialize_from_assets(agent, demand=None, technologies=technologies) assert set(space.dims) == {"asset", "replacement"} assert space.replacement.isin(technologies.technology).all() @@ -235,14 +236,9 @@ def test_init_from_asset(technologies, rng): def test_init_from_asset_no_assets(technologies, rng): - from collections import namedtuple - - from muse.filters import initialize_from_assets - agent = namedtuple("DummyAgent", ["assets"])( xr.Dataset(dict(capacity=xr.DataArray(0))) ) - space = initialize_from_assets(agent, demand=None, technologies=technologies) assert set(space.dims) == {"replacement"} assert space.replacement.isin(technologies.technology).all() From f17f2cdbf0d068ceffdb8aefb1ed9dc43b6157c2 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:45:17 +0100 Subject: [PATCH 19/33] Tidy test_decisions --- tests/test_decisions.py | 125 ++++++++++++++++++++-------------------- 1 file changed, 61 insertions(+), 64 deletions(-) diff --git a/tests/test_decisions.py b/tests/test_decisions.py index a317ba26..6bb9525d 100644 --- a/tests/test_decisions.py +++ b/tests/test_decisions.py @@ -1,11 +1,20 @@ +import numpy as np +from numpy.random import choice, rand from pytest import approx, fixture +from scipy.stats import rankdata +from xarray import Dataset + +from muse.decisions import ( + epsilon_constraints, + lexical_comparison, + retro_epsilon_constraints, + single_objective, + weighted_sum, +) @fixture def objectives(): - from numpy.random import choice - from xarray import Dataset - objectives = Dataset() objectives["replacement"] = choice( list("abcdefghijklmnopqrstuvwxyz"), 10, replace=False @@ -23,20 +32,15 @@ def objectives(): def add_var(coordinates, *dims, factor=100.0): - from numpy.random import rand - shape = tuple(len(coordinates[u]) for u in dims) return dims, (rand(*shape) * factor).astype(type(factor)) def test_weighted_sum(objectives): - from muse.decisions import weighted_sum - weights = {"a": -1, "c": 2} def normalize(objective): - norm = abs(objective).max("replacement") - return objective / norm + return objective / abs(objective).max("replacement") expected = ( normalize(objectives.a) * weights["a"] @@ -50,104 +54,97 @@ def normalize(objective): def test_lexical(): """Test lexical comparison against hand-constructed tuples.""" - from numpy import floor, zeros - from numpy.random import choice, rand - from scipy.stats import rankdata - from xarray import Dataset - - from muse.decisions import lexical_comparison - - a = rand(5, 10) * 10 - 5 - b = rand(5, 10) * 10 - 5 - c = rand(5, 10) * 10 - 5 + shape = (5, 10) + a = rand(*shape) * 10 - 5 + b = rand(*shape) * 10 - 5 + c = rand(*shape) * 10 - 5 parameters = [("b", -rand() * 0.1), ("a", rand() * 0.1), ("c", rand())] - mina = (a * dict(parameters)["a"]).min(1) - minb = -(b * abs(dict(parameters)["b"])).max(1) - minc = (c * dict(parameters)["c"]).min(1) + param_dict = dict(parameters) - expected = zeros(shape=a.shape, dtype=object) - for i in range(expected.shape[0]): - for j in range(expected.shape[1]): + mina = (a * param_dict["a"]).min(1) + minb = -(b * abs(param_dict["b"])).max(1) + minc = (c * param_dict["c"]).min(1) + + expected = np.zeros(shape=shape, dtype=object) + for i in range(shape[0]): + for j in range(shape[1]): expected[i, j] = ( - int(floor(b[i, j] / minb[i])), - int(floor(a[i, j] / mina[i])), + int(np.floor(b[i, j] / minb[i])), + int(np.floor(a[i, j] / mina[i])), c[i, j] / minc[i], ) - objectives = Dataset() - objectives["a"] = ("asset", "replacement"), a - objectives["b"] = ("asset", "replacement"), b - objectives["c"] = ("asset", "replacement"), c - objectives["asset"] = choice( - objectives.replacement, len(objectives.asset), replace=False + objectives = Dataset( + { + "a": (("asset", "replacement"), a), + "b": (("asset", "replacement"), b), + "c": (("asset", "replacement"), c), + } ) + objectives["asset"] = choice(objectives.replacement, shape[0], replace=False) actual = lexical_comparison(objectives, parameters) assert actual.shape == expected.shape - for i in range(expected.shape[0]): + for i in range(shape[0]): assert actual.values[i] == approx(rankdata(expected[i])) def test_epsilon_constraints(objectives): - from numpy import array, isnan - from numpy.random import choice - - from muse.decisions import epsilon_constraints, retro_epsilon_constraints + def reshape_array(size, shape, start=1): + return np.array(range(start, size + start)).reshape(shape) - objectives.b[:] = array(range(1, objectives.b.size + 1)).reshape(objectives.b.shape) - objectives.c[:] = array(range(1, objectives.c.size + 1)).reshape(objectives.c.shape) + objectives.b[:] = reshape_array(objectives.b.size, objectives.b.shape) + objectives.c[:] = reshape_array(objectives.c.size, objectives.c.shape) + # Test case 1: Basic constraints expected = objectives.a * (objectives.asset == objectives.asset) - parameters = [("a", True), ("b", True, objectives.b.max() + 1)] - actual = epsilon_constraints(objectives, parameters) + params = [("a", True), ("b", True, objectives.b.max() + 1)] + actual = epsilon_constraints(objectives, params) assert actual.values == approx(expected.values) + # Test case 2: Negative constraints expected = -objectives.a * (objectives.asset == objectives.asset) - parameters = [("a", False), ("b", True, objectives.b.max() + 1)] - actual = epsilon_constraints(objectives, parameters) + params = [("a", False), ("b", True, objectives.b.max() + 1)] + actual = epsilon_constraints(objectives, params) assert actual.values == approx(expected.values) + # Test case 3: Binary choice constraints objectives.b[:] = choice((1, 2), objectives.b.size).reshape(objectives.b.shape) - parameters = [("a", True), ("b", True, 1.5)] + params = [("a", True), ("b", True, 1.5)] expected = objectives.a.where(objectives.b == 1).fillna(-1) - actual = epsilon_constraints(objectives, parameters, mask=-1) + actual = epsilon_constraints(objectives, params, mask=-1) assert actual.values == approx(expected.values) + # Test case 4: Multiple constraints objectives.b[:] = choice((1, 2), objectives.b.size).reshape(objectives.b.shape) objectives.c[:] = choice((1, 2, 3), objectives.c.size).reshape(objectives.c.shape) - parameters = [("a", True), ("b", True, 1.5), ("c", False, 1.2)] + params = [("a", True), ("b", True, 1.5), ("c", False, 1.2)] condition = (objectives.b == 1) & (objectives.c >= 2).all("other") expected = objectives.a.where(condition).fillna(-1) - actual = epsilon_constraints(objectives, parameters, mask=-1) - assert (isnan(actual) == isnan(expected)).all() + actual = epsilon_constraints(objectives, params, mask=-1) + assert (np.isnan(actual) == np.isnan(expected)).all() assert actual.fillna(0).values == approx(expected.fillna(0).values) - # retro_ should tweak the parameters so that assets are always - # included. Hence expected values cannot be nan. + # Test retro_epsilon_constraints expected = expected.fillna(-1) - while not isnan(expected).any(): + while not np.isnan(expected).any(): objectives.b[:] = choice((1, 2), objectives.b.size).reshape(objectives.b.shape) objectives.c[:] = choice((1, 2, 3), objectives.c.size).reshape( objectives.c.shape ) condition = (objectives.b == 1) & (objectives.c >= 2).all("other") expected = objectives.a.where(condition).min("replacement") - actual = retro_epsilon_constraints(objectives, parameters) - assert not isnan(actual).any() + actual = retro_epsilon_constraints(objectives, params) + assert not np.isnan(actual).any() def test_single_objectives(objectives): - from muse.decisions import single_objective - assert single_objective(objectives, "a").values == approx(objectives.a.values) - actual = single_objective(objectives, ["b", False]) - assert actual.values == approx(-objectives.b.values) + assert single_objective(objectives, ["b", False]).values == approx( + -objectives.b.values + ) -# when developing/debugging, these few lines help setup the input for the -# different tests -if __name__ == "main": - # fmt: off - from tests.agents import test_objectives # noqa - objectives = test_decisions.objectives() # noqa +if __name__ == "__main__": + test_objectives = objectives() # For debugging purposes From a7e93db7b817aaf1815b4df25758c65137c1b342 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:47:06 +0100 Subject: [PATCH 20/33] Tidy test_demand_share --- tests/test_demand_share.py | 337 ++++++++++++++++++++----------------- 1 file changed, 181 insertions(+), 156 deletions(-) diff --git a/tests/test_demand_share.py b/tests/test_demand_share.py index c061d58c..ad4be6e8 100644 --- a/tests/test_demand_share.py +++ b/tests/test_demand_share.py @@ -1,13 +1,70 @@ +from dataclasses import dataclass +from uuid import UUID, uuid4 + import xarray as xr from pytest import approx, fixture, raises +from muse.commodities import is_enduse +from muse.quantities import maximum_production from muse.timeslices import drop_timeslice -from muse.utilities import interpolate_capacity +from muse.utilities import broadcast_over_assets, interpolate_capacity CURRENT_YEAR = 2010 INVESTMENT_YEAR = 2015 +@dataclass +class Agent: + """Test agent with required attributes.""" + + assets: xr.Dataset + category: str = "" + uuid: UUID = None + name: str = "" + region: str = "" + quantity: float = 0.0 + + +def create_test_agents(usa_stock, asia_stock=None, with_new=True): + """Helper to create test agents with standard configuration.""" + agents = [ + Agent(0.3 * usa_stock, "retrofit", uuid4(), "a", "USA", 0.3), + Agent(0.7 * usa_stock, "retrofit", uuid4(), "b", "USA", 0.7), + ] + if with_new: + agents.extend( + [ + Agent(0.0 * usa_stock, "new", uuid4(), "a", "USA", 0.0), + Agent(0.0 * usa_stock, "new", uuid4(), "b", "USA", 0.0), + ] + ) + if asia_stock is not None: + agents.extend( + [ + Agent(asia_stock, "retrofit", uuid4(), "a", "ASEAN", 1.0), + Agent(0 * asia_stock, "new", uuid4(), "a", "ASEAN", 0.0) + if with_new + else None, + ] + ) + return [a for a in agents if a is not None] + + +def create_regional_market(technologies, stock): + """Create market data for given regions.""" + asia_stock = stock.where(stock.region == "ASEAN", drop=True) + usa_stock = stock.where(stock.region == "USA", drop=True) + + asia_market = _matching_market( + broadcast_over_assets(technologies, asia_stock), asia_stock.capacity + ) + usa_market = _matching_market( + broadcast_over_assets(technologies, usa_stock), usa_stock.capacity + ) + market = xr.concat((asia_market, usa_market), dim="region") + return market, asia_stock, usa_stock + + @fixture def _capacity(stock): return interpolate_capacity(stock.capacity, year=[CURRENT_YEAR, INVESTMENT_YEAR]) @@ -22,8 +79,6 @@ def _technologies(technologies, _capacity): @fixture def _market(_technologies, _capacity, timeslice): """A market which matches stocks exactly.""" - from muse.utilities import broadcast_over_assets - _technologies = broadcast_over_assets(_technologies, _capacity) return _matching_market(_technologies, _capacity).transpose( "timeslice", "region", "commodity", "year" @@ -34,16 +89,16 @@ def _matching_market(technologies, capacity): """A market which matches stocks exactly.""" from numpy.random import random - from muse.quantities import consumption, maximum_production + from muse.quantities import consumption as calc_consumption market = xr.Dataset() production = maximum_production(technologies, capacity) - consumption = consumption(technologies, production) + cons = calc_consumption(technologies, production) if "region" in production.coords: production = production.groupby("region") - consumption = consumption.groupby("region") + cons = cons.groupby("region") market["supply"] = production.sum("asset") - market["consumption"] = drop_timeslice(consumption.sum("asset") + market.supply) + market["consumption"] = drop_timeslice(cons.sum("asset") + market.supply) market["prices"] = market.supply.dims, random(market.supply.shape) return market @@ -55,14 +110,64 @@ def test_fixtures(_capacity, _market, _technologies): def test_new_retro_split_zero_unmet(_capacity, _market, _technologies): + """Test that new and retro demands are zero when demand is fully met.""" from muse.demand_share import new_and_retro_demands - from muse.utilities import broadcast_over_assets _technologies = broadcast_over_assets(_technologies, _capacity) share = new_and_retro_demands(_capacity, _market.consumption, _technologies) assert (share == 0).all() +def test_new_retro_split_scenarios(_capacity, _market, _technologies): + """Test various scenarios for new and retro demand splits.""" + from muse.demand_share import new_and_retro_demands + + _technologies = broadcast_over_assets(_technologies, _capacity) + + def check_share_values(share, expect_new_zero=True, expect_retrofit_nonzero=True): + """Helper to check share values under different scenarios.""" + assert (share.new == 0).all() if expect_new_zero else (share.new != 0).any() + assert ( + (share.retrofit != 0).any() + if expect_retrofit_nonzero + else (share.retrofit == 0).all() + ) + + # Test with same consumption in investment year + _market.consumption.loc[{"year": INVESTMENT_YEAR}] = _market.consumption.sel( + year=CURRENT_YEAR + ) + share = new_and_retro_demands(_capacity, _market.consumption, _technologies) + assert (share == 0).all() + + # Test with reduced future capacity + future_unmet = _capacity.copy() + future_unmet.loc[{"year": INVESTMENT_YEAR}] = 0.5 * future_unmet.sel( + year=CURRENT_YEAR + ) + share = new_and_retro_demands(future_unmet, _market.consumption, _technologies) + check_share_values(share) + + # Test with reduced current capacity + current_unmet = _capacity.copy() + current_unmet.loc[{"year": CURRENT_YEAR}] = 0.5 * future_unmet.sel( + year=CURRENT_YEAR + ) + share = new_and_retro_demands(current_unmet, _market.consumption, _technologies) + check_share_values(share) + + # Test with overall reduced capacity + share = new_and_retro_demands(0.5 * _capacity, _market.consumption, _technologies) + check_share_values(share) + + # Test with market supply matching consumption + _market.consumption.loc[{"year": INVESTMENT_YEAR}] = _market.supply.sel( + year=CURRENT_YEAR, drop=True + ).transpose(*_market.consumption.loc[{"year": INVESTMENT_YEAR}].dims) + share = new_and_retro_demands(_capacity, _market.consumption, _technologies) + assert (share == 0).all() + + def test_new_retro_split_zero_consumption_increase(_capacity, _market, _technologies): from muse.demand_share import new_and_retro_demands from muse.utilities import broadcast_over_assets @@ -135,7 +240,6 @@ def test_new_retro_split_zero_new_unmet(_capacity, _market, _technologies): def test_new_retro_accounting_identity(_capacity, _market, _technologies): from muse.demand_share import new_and_retro_demands - from muse.quantities import maximum_production from muse.utilities import broadcast_over_assets _technologies = broadcast_over_assets(_technologies, _capacity) @@ -164,127 +268,65 @@ def test_new_retro_accounting_identity(_capacity, _market, _technologies): assert accounting.values == approx(consumption.values) -def test_demand_split(_capacity, _market, _technologies): - from muse.commodities import is_enduse +def test_demand_split_scenarios(_capacity, _market, _technologies): + """Test demand split scenarios with different agent configurations.""" from muse.demand_share import _inner_split as inner_split - from muse.utilities import broadcast_over_assets + from muse.demand_share import decommissioning_demand - def method(capacity, technologies): - from muse.demand_share import decommissioning_demand + def get_test_demand(): + """Get test demand data for USA region.""" + return _market.consumption.sel( + year=INVESTMENT_YEAR, region="USA", drop=True + ).where(is_enduse(_technologies.comm_usage.sel(commodity=_market.commodity))) - return decommissioning_demand( - technologies, - capacity, - ) + def check_share_results(share, agents_data, expected_shares): + """Verify share results match expectations.""" + enduse = is_enduse(_technologies.comm_usage) + for agent_name in agents_data.keys(): + assert (share[agent_name].sel(commodity=~enduse) == 0).all() - demand = _market.consumption.sel( - year=INVESTMENT_YEAR, region="USA", drop=True - ).where(is_enduse(_technologies.comm_usage.sel(commodity=_market.commodity))) - agents = dict(scully=_capacity, mulder=_capacity) + total = sum(share.values()).sum("asset") + demand = get_test_demand().where(enduse, 0) + demand, total = xr.broadcast(demand, total) + assert demand.values == approx(total.values) + + for agent_name, expected_share in expected_shares.items(): + expected, actual = xr.broadcast(demand, share[agent_name].sum("asset")) + assert actual.values == approx(expected_share * expected.values) + + # Test normal demand split _technologies = broadcast_over_assets(_technologies, _capacity) + agents = dict(scully=_capacity, mulder=_capacity) technodata = dict(scully=_technologies, mulder=_technologies) quantity = dict(scully=("scully", "USA", 0.3), mulder=("mulder", "USA", 0.7)) - share = inner_split(agents, technodata, demand, method, quantity) - - enduse = is_enduse(_technologies.comm_usage) - assert (share["scully"].sel(commodity=~enduse) == 0).all() - assert (share["mulder"].sel(commodity=~enduse) == 0).all() - - total = (share["scully"] + share["mulder"]).sum("asset") - demand = demand.where(enduse, 0) - demand, total = xr.broadcast(demand, total) - assert demand.values == approx(total.values) - expected, actual = xr.broadcast(demand, share["scully"].sum("asset")) - assert actual.values == approx(0.3 * expected.values) - expected, actual = xr.broadcast(demand, share["mulder"].sum("asset")) - assert actual.values == approx(0.7 * expected.values) - -def test_demand_split_zero_share(_capacity, _market, _technologies): - """See issue SgiModel/StarMuse#688.""" - from muse.commodities import is_enduse - from muse.demand_share import _inner_split as inner_split - from muse.utilities import broadcast_over_assets - - def method(capacity, technologies): - from muse.demand_share import decommissioning_demand - - return 0 * decommissioning_demand( - technologies, - capacity, - ) + share = inner_split( + agents, technodata, get_test_demand(), decommissioning_demand, quantity + ) + check_share_results(share, agents, {"scully": 0.3, "mulder": 0.7}) - demand = _market.consumption.sel( - year=INVESTMENT_YEAR, region="USA", drop=True - ).where(is_enduse(_technologies.comm_usage.sel(commodity=_market.commodity))) + # Test zero share scenario agents = dict(scully=0.3 * _capacity, mulder=0.7 * _capacity) - _technologies = broadcast_over_assets(_technologies, _capacity) - technodata = dict(scully=_technologies, mulder=_technologies) quantity = dict(scully=("scully", "USA", 1), mulder=("mulder", "USA", 1)) - share = inner_split(agents, technodata, demand, method, quantity) - enduse = is_enduse(_technologies.comm_usage) - assert (share["scully"].sel(commodity=~enduse) == 0).all() - assert (share["mulder"].sel(commodity=~enduse) == 0).all() + def zero_decom(technologies, capacity): + """Return zero decommissioning demand.""" + return 0 * decommissioning_demand(technologies=technologies, capacity=capacity) - total = (share["scully"] + share["mulder"]).sum("asset") - demand = demand.where(enduse, 0) - demand, total = xr.broadcast(demand, total) - - assert demand.values == approx(total.values, abs=1e-10) - expected, actual = xr.broadcast(demand, share["scully"].sum("asset")) - - assert actual.values == approx(0.5 * expected.values) - expected, actual = xr.broadcast(demand, share["mulder"].sum("asset")) - assert actual.values == approx(0.5 * expected.values) + share = inner_split(agents, technodata, get_test_demand(), zero_decom, quantity) + check_share_results(share, agents, {"scully": 0.5, "mulder": 0.5}) def test_new_retro_demand_share(_technologies, market, timeslice, stock): - from dataclasses import dataclass - from uuid import UUID, uuid4 - - from muse.commodities import is_enduse + """Test new and retro demand share calculations.""" from muse.demand_share import new_and_retro - from muse.utilities import broadcast_over_assets - - asia_stock = stock.where(stock.region == "ASEAN", drop=True) - usa_stock = stock.where(stock.region == "USA", drop=True) - asia_market = _matching_market( - broadcast_over_assets(_technologies, asia_stock), asia_stock.capacity - ) - usa_market = _matching_market( - broadcast_over_assets(_technologies, usa_stock), usa_stock.capacity - ) - market = xr.concat((asia_market, usa_market), dim="region") + market, asia_stock, usa_stock = create_regional_market(_technologies, stock) market.consumption.loc[{"year": 2030}] *= 2 - - # spoof some agents - @dataclass - class Agent: - assets: xr.Dataset - category: str - uuid: UUID - name: str - region: str - quantity: float - - agents = [ - Agent(0.3 * usa_stock, "retrofit", uuid4(), "a", "USA", 0.3), - Agent(0.0 * usa_stock, "new", uuid4(), "a", "USA", 0.0), - Agent(0.7 * usa_stock, "retrofit", uuid4(), "b", "USA", 0.7), - Agent(0.0 * usa_stock, "new", uuid4(), "b", "USA", 0.0), - Agent(asia_stock, "retrofit", uuid4(), "a", "ASEAN", 1.0), - Agent(0 * asia_stock, "new", uuid4(), "a", "ASEAN", 0.0), - ] - + agents = create_test_agents(usa_stock, asia_stock) results = new_and_retro(agents, market.consumption, _technologies) - for _, share in results.groupby("agent"): - assert share.sel( - commodity=~is_enduse(_technologies.comm_usage) - ).values == approx(0) - + # Verify results for each agent uuid_to_category = {agent.uuid: agent.category for agent in agents} uuid_to_name = {agent.uuid: agent.name for agent in agents} for category in {"retrofit", "new"}: @@ -293,63 +335,39 @@ class Agent: for uuid, share in results.groupby("agent") if uuid_to_category[uuid] == category and (share.region == "USA").all() } - expected, actual = xr.broadcast(0.3 * sum(subset.values()), subset["a"]) - assert actual.values == approx(expected.values) + if subset: + expected, actual = xr.broadcast(0.3 * sum(subset.values()), subset["a"]) + assert actual.values == approx(expected.values) def test_standard_demand_share(_technologies, timeslice, stock): - from dataclasses import dataclass - from uuid import UUID, uuid4 - - from muse.commodities import is_enduse + """Test standard demand share calculations.""" from muse.demand_share import standard_demand from muse.errors import RetrofitAgentInStandardDemandShare - from muse.utilities import broadcast_over_assets - - asia_stock = stock.where(stock.region == "ASEAN", drop=True) - usa_stock = stock.where(stock.region == "USA", drop=True) - asia_market = _matching_market( - broadcast_over_assets(_technologies, asia_stock), asia_stock.capacity - ) - usa_market = _matching_market( - broadcast_over_assets(_technologies, usa_stock), usa_stock.capacity - ) - market = xr.concat((asia_market, usa_market), dim="region") + market, asia_stock, usa_stock = create_regional_market(_technologies, stock) market.consumption.loc[{"year": 2030}] *= 2 - # spoof some agents - @dataclass - class Agent: - assets: xr.Dataset - category: str - uuid: UUID - name: str - region: str - quantity: float + # Test that retrofit agents raise error + with raises(RetrofitAgentInStandardDemandShare): + standard_demand( + create_test_agents(usa_stock, asia_stock), market.consumption, _technologies + ) + # Test with only new agents agents = [ - Agent(0.3 * usa_stock, "retrofit", uuid4(), "a", "USA", 0.3), - Agent(0.0 * usa_stock, "new", uuid4(), "a", "USA", 0.0), - Agent(0.7 * usa_stock, "retrofit", uuid4(), "b", "USA", 0.7), - Agent(0.0 * usa_stock, "new", uuid4(), "b", "USA", 0.0), - Agent(asia_stock, "retrofit", uuid4(), "a", "ASEAN", 1.0), - Agent(0 * asia_stock, "new", uuid4(), "a", "ASEAN", 0.0), + Agent(0.3 * usa_stock, "new", uuid4(), "a", "USA", 0.3), + Agent(0.7 * usa_stock, "new", uuid4(), "b", "USA", 0.7), + Agent(asia_stock, "new", uuid4(), "a", "ASEAN", 1.0), ] - - with raises(RetrofitAgentInStandardDemandShare): - standard_demand(agents, market.consumption, _technologies) - - agents = [a for a in agents if a.category != "retrofit"] - results = standard_demand(agents, market.consumption, _technologies) - uuid_to_category = {agent.uuid: agent.category for agent in agents} + # Verify results uuid_to_name = {agent.uuid: agent.name for agent in agents} subset = { uuid_to_name[uuid]: share.sel(commodity=is_enduse(_technologies.comm_usage)) for uuid, share in results.groupby("agent") - if uuid_to_category[uuid] == "new" and (share.region == "USA").all() + if (share.region == "USA").all() } expected, actual = xr.broadcast(0.3 * sum(subset.values()), subset["a"]) assert actual.values == approx(expected.values) @@ -418,18 +436,25 @@ class Agent: def test_decommissioning_demand(_technologies, _capacity, timeslice): - from muse.commodities import is_enduse + """Test decommissioning demand calculations.""" from muse.demand_share import decommissioning_demand - from muse.utilities import broadcast_over_assets _technologies = broadcast_over_assets(_technologies, _capacity) - _capacity.loc[{"year": CURRENT_YEAR}] = current = 1.3 - _capacity.loc[{"year": INVESTMENT_YEAR}] = forecast = 1.0 - _technologies.fixed_outputs[:] = fouts = 0.5 - _technologies.utilization_factor[:] = ufac = 0.4 + # Set test values + current, forecast = 1.3, 1.0 + fixed_outputs, utilization = 0.5, 0.4 + + _capacity.loc[{"year": CURRENT_YEAR}] = current + _capacity.loc[{"year": INVESTMENT_YEAR}] = forecast + _technologies.fixed_outputs[:] = fixed_outputs + _technologies.utilization_factor[:] = utilization + + # Calculate and verify decommissioning demand decom = decommissioning_demand(_technologies, _capacity) assert set(decom.dims) == {"asset", "commodity", "timeslice"} + + expected_decom = utilization * fixed_outputs * (current - forecast) assert decom.sel(commodity=is_enduse(_technologies.comm_usage)).sum( "timeslice" - ).values == approx(ufac * fouts * (current - forecast)) + ).values == approx(expected_decom) From 7413904fa3e6c7960767486d0c8bdbd00bf895c4 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:50:07 +0100 Subject: [PATCH 21/33] Tidy test_constraints --- tests/test_constraints.py | 202 +++++++++++++++----------------------- 1 file changed, 79 insertions(+), 123 deletions(-) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index 8f8b83fc..a55c232a 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -113,12 +113,7 @@ def demand_constraint(market_demand, capacity, search_space, technologies): def max_capacity_expansion(market_demand, capacity, search_space, technologies): from muse.constraints import max_capacity_expansion - return max_capacity_expansion( - market_demand, - capacity, - search_space, - technologies, - ) + return max_capacity_expansion(market_demand, capacity, search_space, technologies) @fixture @@ -156,24 +151,10 @@ def test_constraints_dimensions( max_production, demand_constraint, demand_limiting_capacity, max_capacity_expansion ): # Max production constraint - assert set(max_production.capacity.dims) == { - "asset", - "commodity", - "replacement", - "timeslice", - } - assert set(max_production.production.dims) == { - "asset", - "commodity", - "replacement", - "timeslice", - } - assert set(max_production.b.dims) == { - "asset", - "commodity", - "replacement", - "timeslice", - } + max_prod_dims = {"asset", "commodity", "replacement", "timeslice"} + assert set(max_production.capacity.dims) == max_prod_dims + assert set(max_production.production.dims) == max_prod_dims + assert set(max_production.b.dims) == max_prod_dims # Demand constraint assert set(demand_constraint.capacity.dims) == set() @@ -196,37 +177,25 @@ def test_constraints_dimensions( def test_lp_constraints_matrix_b_is_scalar(constraint, lpcosts): - """B is a scalar. - - When ``b`` is a scalar, the output should be equivalent to a single row matrix, or a - single vector with only decision variables. - """ + """B is a scalar - output should be equivalent to a single row matrix.""" from muse.constraints import lp_constraint_matrix - lpconstraint = lp_constraint_matrix( - xr.DataArray(1), constraint.capacity, lpcosts.capacity - ) - assert lpconstraint.values == approx(-1) - assert set(lpconstraint.dims) == {f"d({x})" for x in lpcosts.capacity.dims} - - lpconstraint = lp_constraint_matrix( - xr.DataArray(1), constraint.production, lpcosts.production - ) - assert lpconstraint.values == approx(1) - assert set(lpconstraint.dims) == {f"d({x})" for x in lpcosts.production.dims} + for attr in ["capacity", "production"]: + lpconstraint = lp_constraint_matrix( + xr.DataArray(1), getattr(constraint, attr), getattr(lpcosts, attr) + ) + expected_value = -1 if attr == "capacity" else 1 + assert lpconstraint.values == approx(expected_value) + assert set(lpconstraint.dims) == { + f"d({x})" for x in getattr(lpcosts, attr).dims + } def test_max_production_constraint_diagonal(constraint, lpcosts): - """Production side of max capacity production is diagonal. - - The production for each timeslice, region, asset, and replacement technology should - not outstrip the assigned for the asset and replacement technology. Hence, the - production side of the constraint is the identity with a -1 factor. The capacity - side is diagonal, but the values reflect the max-production for each timeslices, - commodity and technology. - """ + """Test production side of max capacity production is diagonal.""" from muse.constraints import lp_constraint_matrix + # Test capacity constraints result = lp_constraint_matrix(constraint.b, constraint.capacity, lpcosts.capacity) decision_dims = {f"d({x})" for x in lpcosts.capacity.dims} constraint_dims = { @@ -234,11 +203,14 @@ def test_max_production_constraint_diagonal(constraint, lpcosts): } assert set(result.dims) == decision_dims.union(constraint_dims) + # Test production constraints result = lp_constraint_matrix( constraint.b, constraint.production, lpcosts.production ) decision_dims = {f"d({x})" for x in lpcosts.production.dims} assert set(result.dims) == decision_dims.union(constraint_dims) + + # Verify diagonal matrix result = result.reset_index("d(timeslice)", drop=True).assign_coords( {"d(timeslice)": result["d(timeslice)"].values} ) @@ -251,12 +223,15 @@ def test_lp_constraint(constraint, lpcosts): from muse.constraints import lp_constraint result = lp_constraint(constraint, lpcosts) - decision_dims = {f"d({x})" for x in lpcosts.capacity.dims} constraint_dims = { f"c({x})" for x in set(lpcosts.production.dims).union(constraint.b.dims) } + + # Test capacity constraints + decision_dims = {f"d({x})" for x in lpcosts.capacity.dims} assert set(result.capacity.dims) == decision_dims.union(constraint_dims) + # Test production constraints decision_dims = {f"d({x})" for x in lpcosts.production.dims} assert set(result.production.dims) == decision_dims.union(constraint_dims) stacked = result.production.stack( @@ -270,6 +245,7 @@ def test_lp_constraint(constraint, lpcosts): def test_to_scipy_adapter_maxprod(costs, max_production, commodities, lpcosts): + """Test scipy adapter with max production constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( @@ -305,6 +281,7 @@ def test_to_scipy_adapter_maxprod(costs, max_production, commodities, lpcosts): def test_to_scipy_adapter_demand(costs, demand_constraint, commodities, lpcosts): + """Test scipy adapter with demand constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( @@ -336,6 +313,7 @@ def test_to_scipy_adapter_demand(costs, demand_constraint, commodities, lpcosts) def test_to_scipy_adapter_max_capacity_expansion( costs, max_capacity_expansion, commodities, lpcosts ): + """Test scipy adapter with max capacity expansion constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( @@ -364,27 +342,25 @@ def test_to_scipy_adapter_max_capacity_expansion( assert set(adapter.A_ub[:, :capsize].flatten()) == {0.0, 1.0} -def test_to_scipy_adapter_no_constraint(costs, commodities, lpcosts): +def test_scipy_adapter_no_constraint(costs, commodities, lpcosts): from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory(costs, constraints=[], commodities=commodities) assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) - assert adapter.A_ub is None - assert adapter.b_ub is None - assert adapter.A_eq is None - assert adapter.b_eq is None + assert all( + getattr(adapter, attr) is None for attr in ["A_ub", "b_ub", "A_eq", "b_eq"] + ) assert adapter.c.ndim == 1 - - capsize = lpcosts.capacity.size # number of capacity decision variables - prodsize = lpcosts.production.size # number of production decision variables - assert adapter.c.size == capsize + prodsize + assert adapter.c.size == lpcosts.capacity.size + lpcosts.production.size -def test_back_to_muse_capacity(lpcosts): +def test_back_to_muse_quantities(lpcosts): from muse.constraints import ScipyAdapter data = ScipyAdapter._unified_dataset(lpcosts) + + # Test capacity lpquantity = ScipyAdapter._selected_quantity(data, "capacity") assert set(lpquantity.dims) == {"d(asset)", "d(replacement)"} copy = ScipyAdapter._back_to_muse_quantity( @@ -392,11 +368,7 @@ def test_back_to_muse_capacity(lpcosts): ) assert (copy == lpcosts.capacity).all() - -def test_back_to_muse_production(lpcosts): - from muse.constraints import ScipyAdapter - - data = ScipyAdapter._unified_dataset(lpcosts) + # Test production lpquantity = ScipyAdapter._selected_quantity(data, "production") assert set(lpquantity.dims) == { "d(asset)", @@ -437,6 +409,7 @@ def test_back_to_muse_all(lpcosts): def test_scipy_adapter_back_to_muse(costs, constraints, commodities, lpcosts): + """Test converting back from scipy adapter format to MUSE format.""" from muse.constraints import ScipyAdapter data = ScipyAdapter._unified_dataset(lpcosts) @@ -474,10 +447,11 @@ def test_scipy_adapter_standard_constraints(costs, constraints, commodities): from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory(costs, constraints, commodities=commodities) - maxprod = next(cs for cs in constraints if cs.name == "max_production") - maxcapa = next(cs for cs in constraints if cs.name == "max capacity expansion") - demand = next(cs for cs in constraints if cs.name == "demand") - dlc = next(cs for cs in constraints if cs.name == "demand_limiting_capacity") + constraint_map = {cs.name: cs for cs in constraints} + maxprod = constraint_map["max_production"] + maxcapa = constraint_map["max capacity expansion"] + demand = constraint_map["demand"] + dlc = constraint_map["demand_limiting_capacity"] n_constraints = adapter.b_ub.size n_decision_vars = adapter.c.size @@ -486,10 +460,11 @@ def test_scipy_adapter_standard_constraints(costs, constraints, commodities): assert adapter.b_eq is None assert adapter.A_eq is None assert adapter.A_ub.shape == (n_constraints, n_decision_vars) - assert n_constraints == demand.b.size + maxprod.b.size + maxcapa.b.size + dlc.b.size + assert n_constraints == sum(c.b.size for c in [demand, maxprod, maxcapa, dlc]) def test_scipy_solver(technologies, costs, constraints, commodities): + """Test the scipy solver for demand matching.""" from muse.investments import scipy_match_demand solution = scipy_match_demand( @@ -508,28 +483,19 @@ def test_minimum_service( ): from muse.constraints import minimum_service - minimum_service_constraint = minimum_service( - market_demand, capacity, search_space, technologies - ) - - # test it is none (when appropriate) - assert minimum_service_constraint is None + # Test with no minimum service factor + assert minimum_service(market_demand, capacity, search_space, technologies) is None - # add the column to technologies - minimum_service_factor = 0.4 * xr.ones_like(technologies.technology, dtype=float) - technologies["minimum_service_factor"] = minimum_service_factor - - # append minimum_service_constraint to constraints - minimum_service_constraint = minimum_service( - market_demand, capacity, search_space, technologies + # Test with minimum service factor + technologies["minimum_service_factor"] = 0.4 * xr.ones_like( + technologies.technology, dtype=float ) - constraints.append(minimum_service_constraint) - - # test that it is no longer none - assert isinstance(minimum_service_constraint, xr.Dataset) + min_service = minimum_service(market_demand, capacity, search_space, technologies) + constraints.append(min_service) + assert isinstance(min_service, xr.Dataset) -def test_max_capacity_expansion(max_capacity_expansion): +def test_max_capacity_expansion_properties(max_capacity_expansion): assert (max_capacity_expansion.capacity == 1).all() assert max_capacity_expansion.production == 0 assert max_capacity_expansion.b.dims == ("replacement",) @@ -545,17 +511,10 @@ def test_max_capacity_expansion_no_limits( ): from muse.constraints import max_capacity_expansion - # Without growth limits, the constraint should return None techs = technologies.drop_vars( ["max_capacity_addition", "max_capacity_growth", "total_capacity_limit"] ) - result = max_capacity_expansion( - market_demand, - capacity, - search_space, - techs, - ) - assert result is None + assert max_capacity_expansion(market_demand, capacity, search_space, techs) is None def test_max_capacity_expansion_seed( @@ -566,29 +525,19 @@ def test_max_capacity_expansion_seed( seed = 10 technologies["growth_seed"] = seed - # Scenario 1: Zero capacity - capacity[:] = 0 - result1 = max_capacity_expansion( - market_demand, capacity, search_space, technologies - ) - - # Scenario 2: Existing capacity = seed - capacity.sel(year=2020)[:] = seed - result2 = max_capacity_expansion( - market_demand, capacity, search_space, technologies - ) - - # Scenario 3: Existing capacity > seed - capacity.sel(year=2020)[:] = 2 * seed - result3 = max_capacity_expansion( - market_demand, capacity, search_space, technologies - ) - - # Result with zero capacity should match result with capacity = seed - assert result1.b.values == approx(result2.b.values) + # Test different capacity scenarios + scenarios = [0, seed, 2 * seed] + results = [] + for cap in scenarios: + capacity.sel(year=2020)[:] = cap + results.append( + max_capacity_expansion(market_demand, capacity, search_space, technologies) + ) - # Result with capacity > should not match - assert result1.b.values != approx(result3.b.values) + # Zero capacity should match seed capacity + assert results[0].b.values == approx(results[1].b.values) + # Higher capacity should differ + assert results[0].b.values != approx(results[2].b.values) def test_max_capacity_expansion_infinite_limits( @@ -596,10 +545,12 @@ def test_max_capacity_expansion_infinite_limits( ): from muse.constraints import max_capacity_expansion - # If all limits are infinite, a ValueError should be raised - technologies["max_capacity_addition"] = np.inf - technologies["max_capacity_growth"] = np.inf - technologies["total_capacity_limit"] = np.inf + for limit in [ + "max_capacity_addition", + "max_capacity_growth", + "total_capacity_limit", + ]: + technologies[limit] = np.inf with raises(ValueError): max_capacity_expansion(market_demand, capacity, search_space, technologies) @@ -611,14 +562,19 @@ def test_max_production(max_production): def test_demand_limiting_capacity( demand_limiting_capacity, max_production, demand_constraint ): - assert demand_limiting_capacity.capacity.values == approx( + # Test capacity values + expected_capacity = ( -max_production.capacity.max("timeslice").values if "timeslice" in max_production.capacity.dims else -max_production.capacity.values ) + assert demand_limiting_capacity.capacity.values == approx(expected_capacity) + + # Test production and b values assert demand_limiting_capacity.production == 0 - assert demand_limiting_capacity.b.values == approx( + expected_b = ( demand_constraint.b.max("timeslice").values if "timeslice" in demand_constraint.b.dims else demand_constraint.b.values ) + assert demand_limiting_capacity.b.values == approx(expected_b) From 7bd3b8f964d65ed25de6ce4bee29f757ee5b6cd4 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 13:24:27 +0100 Subject: [PATCH 22/33] Further tidy test_demand_share --- tests/test_demand_share.py | 211 +++++++++++++++++++++++++------------ 1 file changed, 141 insertions(+), 70 deletions(-) diff --git a/tests/test_demand_share.py b/tests/test_demand_share.py index ad4be6e8..36b3e55d 100644 --- a/tests/test_demand_share.py +++ b/tests/test_demand_share.py @@ -26,7 +26,16 @@ class Agent: def create_test_agents(usa_stock, asia_stock=None, with_new=True): - """Helper to create test agents with standard configuration.""" + """Helper to create test agents with standard configuration. + + Args: + usa_stock: Stock data for USA region + asia_stock: Optional stock data for ASEAN region + with_new: Whether to include new capacity agents + + Returns: + List of Agent objects with specified configurations + """ agents = [ Agent(0.3 * usa_stock, "retrofit", uuid4(), "a", "USA", 0.3), Agent(0.7 * usa_stock, "retrofit", uuid4(), "b", "USA", 0.7), @@ -51,7 +60,15 @@ def create_test_agents(usa_stock, asia_stock=None, with_new=True): def create_regional_market(technologies, stock): - """Create market data for given regions.""" + """Create market data for given regions. + + Args: + technologies: Technology parameters + stock: Stock data containing regional information + + Returns: + Tuple of (market data, asia stock subset, usa stock subset) + """ asia_stock = stock.where(stock.region == "ASEAN", drop=True) usa_stock = stock.where(stock.region == "USA", drop=True) @@ -67,18 +84,19 @@ def create_regional_market(technologies, stock): @fixture def _capacity(stock): + """Create interpolated capacity fixture.""" return interpolate_capacity(stock.capacity, year=[CURRENT_YEAR, INVESTMENT_YEAR]) @fixture def _technologies(technologies, _capacity): - """Technology parameters for the sector.""" + """Create technology parameters fixture for the sector.""" return technologies.interp(year=INVESTMENT_YEAR) @fixture def _market(_technologies, _capacity, timeslice): - """A market which matches stocks exactly.""" + """Create a market fixture which matches stocks exactly.""" _technologies = broadcast_over_assets(_technologies, _capacity) return _matching_market(_technologies, _capacity).transpose( "timeslice", "region", "commodity", "year" @@ -86,7 +104,15 @@ def _market(_technologies, _capacity, timeslice): def _matching_market(technologies, capacity): - """A market which matches stocks exactly.""" + """Create a market which matches stocks exactly. + + Args: + technologies: Technology parameters + capacity: Capacity data + + Returns: + Market dataset with supply, consumption and prices + """ from numpy.random import random from muse.quantities import consumption as calc_consumption @@ -103,14 +129,31 @@ def _matching_market(technologies, capacity): return market +def verify_share_values(share, expect_new_zero=True, expect_retrofit_nonzero=True): + """Helper to verify share values under different scenarios. + + Args: + share: Share values to verify + expect_new_zero: Whether new capacity share should be zero + expect_retrofit_nonzero: Whether retrofit share should be non-zero + """ + assert (share.new == 0).all() if expect_new_zero else (share.new != 0).any() + assert ( + (share.retrofit != 0).any() + if expect_retrofit_nonzero + else (share.retrofit == 0).all() + ) + + def test_fixtures(_capacity, _market, _technologies): + """Verify that test fixtures have the expected dimensions.""" assert set(_capacity.dims) == {"asset", "year"} assert set(_market.dims) == {"commodity", "region", "year", "timeslice"} assert set(_technologies.dims) == {"technology", "region", "commodity"} def test_new_retro_split_zero_unmet(_capacity, _market, _technologies): - """Test that new and retro demands are zero when demand is fully met.""" + """Test that new and retrofit demands are zero when demand is fully met.""" from muse.demand_share import new_and_retro_demands _technologies = broadcast_over_assets(_technologies, _capacity) @@ -119,20 +162,19 @@ def test_new_retro_split_zero_unmet(_capacity, _market, _technologies): def test_new_retro_split_scenarios(_capacity, _market, _technologies): - """Test various scenarios for new and retro demand splits.""" + """Test various scenarios for new and retrofit demand splits. + + Tests: + 1. Same consumption in investment year + 2. Reduced future capacity + 3. Reduced current capacity + 4. Overall reduced capacity + 5. Market supply matching consumption + """ from muse.demand_share import new_and_retro_demands _technologies = broadcast_over_assets(_technologies, _capacity) - def check_share_values(share, expect_new_zero=True, expect_retrofit_nonzero=True): - """Helper to check share values under different scenarios.""" - assert (share.new == 0).all() if expect_new_zero else (share.new != 0).any() - assert ( - (share.retrofit != 0).any() - if expect_retrofit_nonzero - else (share.retrofit == 0).all() - ) - # Test with same consumption in investment year _market.consumption.loc[{"year": INVESTMENT_YEAR}] = _market.consumption.sel( year=CURRENT_YEAR @@ -146,7 +188,7 @@ def check_share_values(share, expect_new_zero=True, expect_retrofit_nonzero=True year=CURRENT_YEAR ) share = new_and_retro_demands(future_unmet, _market.consumption, _technologies) - check_share_values(share) + verify_share_values(share) # Test with reduced current capacity current_unmet = _capacity.copy() @@ -154,11 +196,11 @@ def check_share_values(share, expect_new_zero=True, expect_retrofit_nonzero=True year=CURRENT_YEAR ) share = new_and_retro_demands(current_unmet, _market.consumption, _technologies) - check_share_values(share) + verify_share_values(share) # Test with overall reduced capacity share = new_and_retro_demands(0.5 * _capacity, _market.consumption, _technologies) - check_share_values(share) + verify_share_values(share) # Test with market supply matching consumption _market.consumption.loc[{"year": INVESTMENT_YEAR}] = _market.supply.sel( @@ -169,78 +211,80 @@ def check_share_values(share, expect_new_zero=True, expect_retrofit_nonzero=True def test_new_retro_split_zero_consumption_increase(_capacity, _market, _technologies): + """Test new and retrofit demand splits with no consumption increase. + + Tests various capacity scenarios when consumption remains constant. + """ from muse.demand_share import new_and_retro_demands - from muse.utilities import broadcast_over_assets _technologies = broadcast_over_assets(_technologies, _capacity) + # Base case - same consumption _market.consumption.loc[{"year": INVESTMENT_YEAR}] = _market.consumption.sel( year=CURRENT_YEAR ) share = new_and_retro_demands(_capacity, _market.consumption, _technologies) assert (share == 0).all() - future_unmet = _capacity.copy() - future_unmet.loc[{"year": INVESTMENT_YEAR}] = 0.5 * future_unmet.sel( - year=CURRENT_YEAR - ) - share = new_and_retro_demands(future_unmet, _market.consumption, _technologies) - assert (share.new == 0).all() - assert (share.retrofit != 0).any() - - current_unmet = _capacity.copy() - current_unmet.loc[{"year": CURRENT_YEAR}] = 0.5 * future_unmet.sel( - year=CURRENT_YEAR - ) - share = new_and_retro_demands(current_unmet, _market.consumption, _technologies) - assert (share.new == 0).all() - assert (share.retrofit != 0).any() + # Test capacity reduction scenarios + scenarios = [ + ("future", lambda x: x.loc[{"year": INVESTMENT_YEAR}], 0.5), + ("current", lambda x: x.loc[{"year": CURRENT_YEAR}], 0.5), + ("overall", lambda x: x, 0.5), + ] - share = new_and_retro_demands(0.5 * _capacity, _market.consumption, _technologies) - assert (share.new == 0).all() - assert (share.retrofit != 0).any() + for name, selector, factor in scenarios: + modified_capacity = _capacity.copy() + selector(modified_capacity)[:] = factor * modified_capacity.sel( + year=CURRENT_YEAR + ) + share = new_and_retro_demands( + modified_capacity, _market.consumption, _technologies + ) + verify_share_values(share) def test_new_retro_split_zero_new_unmet(_capacity, _market, _technologies): + """Test new and retrofit demand splits with zero new unmet demand. + + Tests that retrofit demand is properly allocated when there is no new unmet demand. + """ from muse.demand_share import new_and_retro_demands - from muse.utilities import broadcast_over_assets _technologies = broadcast_over_assets(_technologies, _capacity) + # Set market consumption to match current supply _market.consumption.loc[{"year": INVESTMENT_YEAR}] = _market.supply.sel( year=CURRENT_YEAR, drop=True ).transpose(*_market.consumption.loc[{"year": INVESTMENT_YEAR}].dims) share = new_and_retro_demands(_capacity, _market.consumption, _technologies) assert (share == 0).all() - future_unmet = _capacity.copy() - future_unmet.loc[{"year": INVESTMENT_YEAR}] = 0.5 * future_unmet.sel( - year=CURRENT_YEAR - ) - share = new_and_retro_demands(future_unmet, _market.consumption, _technologies) - assert (share.new == 0).all() - assert (share.retrofit != 0).any() - - current_unmet = _capacity.copy() - current_unmet.loc[{"year": CURRENT_YEAR}] = 0.5 * future_unmet.sel( - year=CURRENT_YEAR - ) - share = new_and_retro_demands(current_unmet, _market.consumption, _technologies) - assert (share.new == 0).all() - assert (share.retrofit != 0).any() + # Test capacity reduction scenarios + scenarios = [ + ("future", lambda x: x.loc[{"year": INVESTMENT_YEAR}], 0.5), + ("current", lambda x: x.loc[{"year": CURRENT_YEAR}], 0.5), + ("overall", lambda x: x, 0.5), + ] - share = new_and_retro_demands( - 0.5 * _capacity, - _market.consumption, - _technologies, - ) - assert (share.new == 0).all() - assert (share.retrofit != 0).any() + for name, selector, factor in scenarios: + modified_capacity = _capacity.copy() + selector(modified_capacity)[:] = factor * modified_capacity.sel( + year=CURRENT_YEAR + ) + share = new_and_retro_demands( + modified_capacity, _market.consumption, _technologies + ) + verify_share_values(share) def test_new_retro_accounting_identity(_capacity, _market, _technologies): + """Test that new and retrofit demands satisfy accounting identity. + + Verifies that the sum of new and retrofit demands plus serviced demand equals total + demand. + """ from muse.demand_share import new_and_retro_demands - from muse.utilities import broadcast_over_assets _technologies = broadcast_over_assets(_technologies, _capacity) @@ -257,9 +301,12 @@ def test_new_retro_accounting_identity(_capacity, _market, _technologies): ) consumption = _market.consumption.sel(year=INVESTMENT_YEAR) + # Verify accounting identity components assert (share.new > -1e-8).all() assert (share.retrofit > -1e-8).all() assert ((share.new + share.retrofit).where(consumption < serviced, 0) < 1e-8).all() + + # Verify total accounting identity accounting = ( (share.new + share.retrofit + serviced) .where(consumption - serviced > 0, consumption) @@ -269,7 +316,11 @@ def test_new_retro_accounting_identity(_capacity, _market, _technologies): def test_demand_split_scenarios(_capacity, _market, _technologies): - """Test demand split scenarios with different agent configurations.""" + """Test demand split scenarios with different agent configurations. + + Tests demand splitting between agents with different quantities and verifies + that shares are properly allocated. + """ from muse.demand_share import _inner_split as inner_split from muse.demand_share import decommissioning_demand @@ -318,7 +369,11 @@ def zero_decom(technologies, capacity): def test_new_retro_demand_share(_technologies, market, timeslice, stock): - """Test new and retro demand share calculations.""" + """Test new and retrofit demand share calculations. + + Verifies that demand is properly shared between new and retrofit agents + across different regions. + """ from muse.demand_share import new_and_retro market, asia_stock, usa_stock = create_regional_market(_technologies, stock) @@ -341,7 +396,12 @@ def test_new_retro_demand_share(_technologies, market, timeslice, stock): def test_standard_demand_share(_technologies, timeslice, stock): - """Test standard demand share calculations.""" + """Test standard demand share calculations. + + Verifies that: + 1. Retrofit agents raise appropriate error + 2. New agents receive proper demand shares + """ from muse.demand_share import standard_demand from muse.errors import RetrofitAgentInStandardDemandShare @@ -374,6 +434,13 @@ def test_standard_demand_share(_technologies, timeslice, stock): def test_unmet_forecast_demand(_technologies, timeslice, stock): + """Test unmet forecast demand calculations. + + Tests three scenarios: + 1. Fully met demand + 2. Excess capacity + 3. Insufficient capacity + """ from dataclasses import dataclass from muse.commodities import is_enduse @@ -391,12 +458,11 @@ def test_unmet_forecast_demand(_technologies, timeslice, stock): ) market = xr.concat((asia_market, usa_market), dim="region") - # spoof some agents @dataclass class Agent: assets: xr.Dataset - # First ensure that the demand is fully met + # Test fully met demand agents = [ Agent(0.3 * usa_stock), Agent(0.7 * usa_stock), @@ -406,7 +472,7 @@ class Agent: assert set(result.dims) == set(market.consumption.dims) - {"year"} assert result.values == approx(0) - # Then try with too little demand + # Test excess capacity agents = [ Agent(0.4 * usa_stock), Agent(0.8 * usa_stock), @@ -420,7 +486,7 @@ class Agent: assert set(result.dims) == set(market.consumption.dims) - {"year"} assert result.values == approx(0) - # Then try too little capacity + # Test insufficient capacity agents = [ Agent(0.5 * usa_stock), Agent(0.5 * asia_stock), @@ -436,7 +502,12 @@ class Agent: def test_decommissioning_demand(_technologies, _capacity, timeslice): - """Test decommissioning demand calculations.""" + """Test decommissioning demand calculations. + + Verifies that decommissioning demand is correctly calculated based on: + 1. Capacity changes between current and forecast years + 2. Fixed outputs and utilization factors + """ from muse.demand_share import decommissioning_demand _technologies = broadcast_over_assets(_technologies, _capacity) From 8031587e6171613d6e6093528068c5d16a67bca6 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Thu, 5 Jun 2025 17:46:25 +0100 Subject: [PATCH 23/33] Revert "Tidy test_outputs" This reverts commit d324bc3ed1fd7d59722d050814cc88493d8b4103. --- tests/test_outputs.py | 472 ++++++++++++++++++++++-------------------- 1 file changed, 243 insertions(+), 229 deletions(-) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 89b2eb2b..10c0ab9d 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -13,7 +13,6 @@ @fixture def streetcred(save_registries): - """Create a test output quantity that returns random data.""" from muse.outputs.sector import register_output_quantity @register_output_quantity @@ -29,103 +28,83 @@ def streetcred(*args, **kwargs): ) -@fixture -def market(): - """Common market fixture used in multiple tests.""" - return xr.DataArray([1], coords={"year": [2010]}, dims="year") - +@mark.usefixtures("streetcred") +def test_save_with_dir(tmpdir): + from pandas import read_csv -@fixture -def base_config(tmpdir): - """Common config fixture used in multiple tests.""" path = Path(tmpdir) / "results" / "stuff" - return { + config = { "filename": path / "{Sector}{year}{Quantity}.csv", "quantity": "streetcred", } - - -def create_test_data_array(values, coords=None, name="test"): - """Helper function to create test DataArrays.""" - if coords is None: - coords = dict(a=[2, 4]) - return xr.DataArray(values, coords=coords, dims="a", name=name) - - -def assert_file_exists_and_readable(path, expected_columns=None): - """Helper to verify file exists and can be read.""" - assert path.exists() and path.is_file() - df = pd.read_csv(path) - if expected_columns: - assert set(df.columns) == set(expected_columns) - return df - - -@mark.usefixtures("streetcred") -def test_save_with_dir(tmpdir, market, base_config): - """Test saving output to directory with sector and year in filename.""" - result = factory(base_config, sector_name="Yoyo")(market, None, None) + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + # can use None because we **know** none of the arguments are used here + result = factory(config, sector_name="Yoyo")(market, None, None) assert len(result) == 1 - expected_path = Path(base_config["filename"]).parent / "Yoyo2010Streetcred.csv" - assert result[0] == expected_path - assert_file_exists_and_readable(result[0]) + assert result[0] == path / "Yoyo2010Streetcred.csv" + assert result[0].exists() + assert result[0].is_file() + read_csv(result[0]) @mark.usefixtures("streetcred") -def test_overwrite(tmpdir, market, base_config): - """Test file overwrite behavior.""" - outputter = factory(base_config, sector_name="Yoyo") +def test_overwrite(tmpdir): + from pytest import raises + + path = Path(tmpdir) / "results" / "stuff" + config = { + "filename": path / "{Sector}{year}{Quantity}.csv", + "quantity": "streetcred", + } + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + # can use None because we **know** none of the arguments are used here + outputter = factory(config, sector_name="Yoyo") result = outputter(market, None, None) - expected_path = Path(base_config["filename"]).parent / "Yoyo2010Streetcred.csv" - assert result[0] == expected_path - assert_file_exists_and_readable(result[0]) + assert result[0] == path / "Yoyo2010Streetcred.csv" + assert result[0].is_file() # default is to never overwrite with raises(IOError): outputter(market, None, None) - base_config["overwrite"] = True - factory(base_config, sector_name="Yoyo")(market, None, None) + config["overwrite"] = True + factory(config, sector_name="Yoyo")(market, None, None) @mark.usefixtures("streetcred") -@mark.parametrize( - "config_type,suffix", - [ - ("suffix", "nc"), - ("sink", "nc"), - ], -) -def test_save_with_path_to_nc(tmpdir, market, base_config, config_type, suffix): - """Test saving output to NC file with different config types.""" +def test_save_with_path_to_nc_with_suffix(tmpdir): path = Path(tmpdir) / "results" / "stuff" - if config_type == "suffix": - config = { - "filename": path / "{Sector}{year}{Quantity}{suffix}", - "quantity": "streetcred", - "suffix": suffix, - } - else: - config = { - "filename": path / "{sector}{year}{quantity}.csv", - "quantity": "streetcred", - "sink": suffix, - } + config = { + "filename": path / "{Sector}{year}{Quantity}{suffix}", + "quantity": "streetcred", + "suffix": "nc", + } + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + # can use None because we **know** none of the arguments are used here result = factory(config, sector_name="Yoyo")(market, None, None) - expected_path = path / ( - f"{'Yoyo' if config_type == 'suffix' else 'yoyo'}" - f"2010" - f"{'Streetcred' if config_type == 'suffix' else 'streetcred'}" - f".{'nc' if config_type == 'suffix' else 'csv'}" - ) - assert result[0] == expected_path + assert result[0] == path / "Yoyo2010Streetcred.nc" + assert result[0].is_file() + xr.open_dataset(result[0]) + + +@mark.usefixtures("streetcred") +def test_save_with_path_to_nc_with_sink(tmpdir): + path = Path(tmpdir) / "results" / "stuff" + # can use None because we **know** none of the arguments are used here + config = { + "filename": path / "{sector}{year}{quantity}.csv", + "quantity": "streetcred", + "sink": "nc", + } + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + result = factory(config, sector_name="Yoyo")(market, None, None) + assert result[0] == path / "yoyo2010streetcred.csv" assert result[0].is_file() xr.open_dataset(result[0]) @mark.usefixtures("streetcred") -def test_save_with_fullpath_to_excel(tmpdir, market): - """Test saving output to Excel file.""" +def test_save_with_fullpath_to_excel_with_sink(tmpdir): from warnings import simplefilter from pandas import read_excel @@ -135,38 +114,25 @@ def test_save_with_fullpath_to_excel(tmpdir, market): path = Path(tmpdir) / "results" / "stuff" / "this.xlsx" config = {"filename": path, "quantity": "streetcred", "sink": "xlsx"} + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + # can use None because we **know** none of the arguments are used here result = factory(config, sector_name="Yoyo")(market, None, None) assert result[0] == path - assert_file_exists_and_readable(result[0]) + assert result[0].is_file() read_excel(result[0]) -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_functions(mock_consolidate): - """Test output functions (capacity, production, lcoe) with common setup.""" - from muse.outputs.cache import capacity, lcoe, production - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - for func, quantity in [ - (capacity, "capacity"), - (production, "production"), - (lcoe, "lcoe"), - ]: - func(cached, agents) - mock_consolidate.assert_called_once_with(quantity, cached, agents) - mock_consolidate.reset_mock() - - @mark.usefixtures("streetcred") -def test_no_sink_or_suffix(tmpdir, market): - """Test default sink and suffix behavior.""" +def test_no_sink_or_suffix(tmpdir): + from muse.outputs.sector import factory + config = dict( quantity="streetcred", filename=f"{tmpdir}/{{Sector}}{{Quantity}}{{year}}{{suffix}}", ) - result = factory(config)(market, None, None) + outputs = factory(config) + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + result = outputs(market, None, None) assert len(result) == 1 assert result[0].is_file() assert result[0].suffix == ".csv" @@ -174,7 +140,6 @@ def test_no_sink_or_suffix(tmpdir, market): @mark.usefixtures("save_registries") def test_can_register_class(): - """Test class registration functionality.""" from muse.outputs.sinks import factory, register_output_sink @register_output_sink @@ -186,14 +151,12 @@ def __init__(self, sector, some_args=3): def __call__(self, x): pass - # Test default arguments settings = {"sink": {"name": "AClass"}} sink = factory(settings, sector_name="yoyo") assert isinstance(sink, AClass) assert sink.sector == "yoyo" assert sink.some_args == 3 - # Test custom arguments settings = {"sink": {"name": "AClass", "some_args": 5}} sink = factory(settings, sector_name="yoyo") assert isinstance(sink, AClass) @@ -203,7 +166,6 @@ def __call__(self, x): @mark.usefixtures("save_registries") def test_can_register_function(): - """Test function registration functionality.""" from muse.outputs.sinks import factory, register_output_sink @register_output_sink @@ -217,10 +179,8 @@ def a_function(*args): @mark.usefixtures("save_registries") def test_yearly_aggregate(): - """Test yearly aggregation with custom sink.""" from muse.outputs.sinks import factory, register_output_sink - # Setup tracking variables received_data = None gyear = None gsector = None @@ -242,17 +202,16 @@ def dummy(data, year: int, sector: str, overwrite: bool) -> MySpecialReturn: dict(overwrite=True, sink=dict(aggregate="dummy")), sector_name="yoyo" ) - # Test first year - data = create_test_data_array([1, 0], name="nada") + data = xr.DataArray([1, 0], coords=dict(a=[2, 4]), dims="a", name="nada") data["year"] = 2010 + assert isinstance(sink(data, 2010), MySpecialReturn) assert gyear == 2010 assert gsector == "yoyo" assert goverwrite is True assert isinstance(received_data, pd.DataFrame) - # Test second year - data = create_test_data_array([0, 1], name="nada") + data = xr.DataArray([0, 1], coords=dict(a=[2, 4]), dims="a", name="nada") data["year"] = 2020 assert isinstance(sink(data, 2020), MySpecialReturn) assert gyear == 2020 @@ -262,42 +221,34 @@ def dummy(data, year: int, sector: str, overwrite: bool) -> MySpecialReturn: def test_yearly_aggregate_file(tmpdir): - """Test yearly aggregation to file with multiple years of data.""" from muse.outputs.sinks import factory path = Path(tmpdir) / "file.csv" sink = factory(dict(filename=str(path), sink="aggregate"), sector_name="yoyo") - def verify_year_data(values, year, expected_rows): - data = create_test_data_array(values, name="georges") - data["year"] = year - assert sink(data, year) == path - df = assert_file_exists_and_readable(path, {"year", "georges"}) - assert df.shape[0] == expected_rows - return df - - # Test first year - verify_year_data([1, 0], 2010, 2) - - # Test second year (should append to existing file) - df2 = verify_year_data([0, 1], 2020, 4) + data = xr.DataArray([1, 0], coords=dict(a=[2, 4]), dims="a", name="georges") + data["year"] = 2010 + assert sink(data, 2010) == path + dataframe = pd.read_csv(path) + assert set(dataframe.columns) == {"year", "georges"} + assert dataframe.shape[0] == 2 - # Verify data from both years is present - assert set(df2.year.unique()) == {2010, 2020} - assert df2[df2.year == 2010].georges.tolist() == [1, 0] - assert df2[df2.year == 2020].georges.tolist() == [0, 1] + data = xr.DataArray([0, 1], coords=dict(a=[2, 4]), dims="a", name="georges") + data["year"] = 2020 + assert sink(data, 2020) == path + dataframe = pd.read_csv(path) + assert set(dataframe.columns) == {"year", "georges"} + assert dataframe.shape[0] == 4 def test_yearly_aggregate_no_outputs(tmpdir): - """Test behavior with no outputs configured.""" from muse.outputs.mca import factory outputs = factory() assert len(outputs(None, year=2010)) == 0 -def setup_mca_test(tmpdir, outputs_config): - """Helper function to set up MCA tests.""" +def test_mca_aggregate_outputs(tmpdir): from toml import dump, load from muse import examples @@ -305,20 +256,17 @@ def setup_mca_test(tmpdir, outputs_config): examples.copy_model(path=str(tmpdir)) settings = load(str(tmpdir / "model" / "settings.toml")) - settings["outputs"] = [outputs_config] + settings["outputs"] = [ + dict(filename="{path}/{Quantity}{suffix}", quantity="prices", sink="aggregate") + ] settings["time_framework"] = settings["time_framework"][:2] dump(settings, (tmpdir / "model" / "settings.toml")) - return MCA.factory(str(tmpdir / "model" / "settings.toml")) - -def test_mca_aggregate_outputs(tmpdir): - """Test MCA aggregate outputs.""" - mca = setup_mca_test( - tmpdir, - dict(filename="{path}/{Quantity}{suffix}", quantity="prices", sink="aggregate"), - ) + mca = MCA.factory(str(tmpdir / "model" / "settings.toml")) mca.run() + assert (tmpdir / "model" / "Prices.csv").exists() + # TODO: should pass again after #612 # data = pd.read_csv(tmpdir / "model" / "Prices.csv") # assert set(data.year) == set(settings["time_framework"]) @@ -326,10 +274,23 @@ def test_mca_aggregate_outputs(tmpdir): @mark.usefixtures("save_registries") def test_path_formatting(tmpdir): - """Test path formatting with dummy sink and quantity.""" + from toml import dump, load + + from muse.examples import copy_model + from muse.mca import MCA from muse.outputs.mca import register_output_quantity from muse.outputs.sinks import register_output_sink, sink_to_file + # Copy the data to tmpdir + copy_model(path=tmpdir) + + settings_file = tmpdir / "model" / "settings.toml" + settings = load(settings_file) + settings["outputs"] = [ + dict(quantity="dummy", sink="to_dummy", filename="{path}/{Quantity}{suffix}") + ] + dump(settings, (settings_file)) + @register_output_sink(name="dummy_sink") @sink_to_file(".dummy") def to_dummy(quantity, filename, **params) -> None: @@ -339,11 +300,12 @@ def to_dummy(quantity, filename, **params) -> None: def dummy(market, **kwargs): return xr.DataArray() - mca = setup_mca_test( - tmpdir, - dict(quantity="dummy", sink="to_dummy", filename="{path}/{Quantity}{suffix}"), + mca = MCA.factory(Path(settings_file)) + assert mca.outputs(mca.market)[0] == Path( + settings["outputs"][0]["filename"].format( + path=tmpdir / "model", Quantity="Dummy", suffix=".dummy" + ) ) - assert mca.outputs(mca.market)[0] == Path(tmpdir / "model" / "Dummy.dummy") def test_register_output_quantity_cache(): @@ -357,59 +319,64 @@ def dummy_quantity(*args): class TestOutputCache: - @fixture - def output_params(self): - return [dict(quantity="height"), dict(quantity="width")] - - @fixture - def output_quantities(self, output_params): - quantities = {q["quantity"]: lambda _: None for q in output_params} - quantities["depth"] = lambda _: None - return quantities - - @fixture - def topic(self): - return "BBC Muse" - @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_init( - self, mock_factory, mock_subscribe, output_params, output_quantities, topic - ): + def test_init(self, mock_factory, mock_subscribe): from muse.outputs.cache import OutputCache + param = [dict(quantity="height"), dict(quantity="width")] + output_quantities = {q["quantity"]: lambda _: None for q in param} + output_quantities["depth"] = lambda _: None + topic = "BBC Muse" + output_cache = OutputCache( - *output_params, output_quantities=output_quantities, topic=topic + *param, output_quantities=output_quantities, topic=topic ) - assert mock_factory.call_count == len(output_params) + + assert mock_factory.call_count == len(param) mock_subscribe.assert_called_once_with(output_cache.cache, topic) @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_cache( - self, mock_factory, mock_subscribe, output_params, output_quantities, topic - ): + def test_cache(self, mock_factory, mock_subscribe): + import xarray as xr + from muse.outputs.cache import OutputCache + param = [dict(quantity="height"), dict(quantity="width")] + output_quantities = {q["quantity"]: lambda _: None for q in param} + output_quantities["depth"] = lambda _: None + topic = "BBC Muse" + output_cache = OutputCache( - *output_params, output_quantities=output_quantities, topic=topic + *param, output_quantities=output_quantities, topic=topic ) + output_cache.cache(dict(height=xr.DataArray(), depth=xr.DataArray())) + assert len(output_cache.to_save.get("height")) == 1 assert len(output_cache.to_save.get("depth", [])) == 0 @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_consolidate_cache( - self, mock_factory, mock_subscribe, output_params, output_quantities, topic - ): + def test_consolidate_cache(self, mock_factory, mock_subscribe): + import xarray as xr + from muse.outputs.cache import OutputCache + param = [dict(quantity="height"), dict(quantity="width")] + output_quantities = {q["quantity"]: lambda _: None for q in param} + output_quantities["depth"] = lambda _: None + topic = "BBC Muse" + year = 2042 + output_cache = OutputCache( - *output_params, output_quantities=output_quantities, topic=topic + *param, output_quantities=output_quantities, topic=topic ) + output_cache.cache(dict(height=xr.DataArray())) - output_cache.consolidate_cache(2042) + output_cache.consolidate_cache(year) + output_cache.factory["height"].assert_called_once() @@ -421,12 +388,9 @@ def test_cache_quantity(mock_match, mock_send): result = {"mass": 42} mock_match.return_value = result - def verify_message_sent(): - mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) - mock_send.reset_mock() - cache_quantity(**result) - verify_message_sent() + mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) + mock_send.reset_mock() with raises(ValueError): cache_quantity(function=lambda: None, mass=42) @@ -443,7 +407,7 @@ def fun2(): fun2() mock_match.assert_called_once_with("mass", 42) - verify_message_sent() + mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) def test_match_quantities(): @@ -451,27 +415,37 @@ def test_match_quantities(): from muse.outputs.cache import match_quantities + q = "mass" + da = xr.DataArray(name=q) + ds = xr.Dataset({q: da}) + def assert_equal(a: dict[str, xr.DataArray], b: dict[str, xr.DataArray]): assert set(a.keys()) == set(b.keys()) for k in a: xr.testing.assert_equal(a[k], b[k]) - # Test single quantity with DataArray - q = "mass" - da = xr.DataArray(name=q) - ds = xr.Dataset({q: da}) - assert_equal(match_quantities(quantity=q, data=da), {q: da}) - assert_equal(match_quantities(quantity=q, data=ds), {q: da}) + actual = match_quantities(quantity=q, data=da) + assert_equal(actual, {q: da}) + + actual = match_quantities(quantity=q, data=ds) + assert_equal(actual, {q: da}) - # Test multiple quantities with Dataset p = "height" ds = xr.Dataset({q: da, p: da, "rubish": da}) - assert_equal(match_quantities(quantity=[q, p], data=ds), {q: da, p: da}) - assert_equal(match_quantities(quantity=[q, p], data=[da, da]), {q: da, p: da}) + actual = match_quantities(quantity=[q, p], data=ds) + assert_equal(actual, {q: da, p: da}) + + actual = match_quantities(quantity=[q, p], data=[da, da]) + assert_equal(actual, {q: da, p: da}) - # Test error cases with raises(ValueError): - match_quantities(quantity=[q, p], data=[da]) + match_quantities( + quantity=[q, p], + data=[ + da, + ], + ) + with raises(TypeError): match_quantities(quantity=[q, p], data=42) @@ -490,54 +464,51 @@ def test_extract_agents(mock_extract): def test_extract_agents_internal(newcapa_agent, retro_agent): - """Test internal agent extraction.""" from types import SimpleNamespace from muse.outputs.cache import extract_agents_internal - def setup_agent(agent, name): - agent.name = name - return agent - - agents = [setup_agent(newcapa_agent, "A1"), setup_agent(retro_agent, "A2")] - sector = SimpleNamespace(name="IT", agents=agents) + newcapa_agent.name = "A1" + retro_agent.name = "A2" + sector = SimpleNamespace(name="IT", agents=[newcapa_agent, retro_agent]) actual = extract_agents_internal(sector) - expected_keys = ("agent", "category", "sector", "dst_region") - - for agent in agents: + for agent in [newcapa_agent, retro_agent]: assert agent.uuid in actual - assert tuple(actual[agent.uuid].keys()) == expected_keys - agent_data = actual[agent.uuid] - assert agent_data["agent"] == agent.name - assert agent_data["category"] == agent.category - assert agent_data["sector"] == "IT" - assert agent_data["dst_region"] == agent.region + assert tuple(actual[agent.uuid].keys()) == ( + "agent", + "category", + "sector", + "dst_region", + ) + assert actual[agent.uuid]["agent"] == agent.name + assert actual[agent.uuid]["category"] == agent.category + assert actual[agent.uuid]["sector"] == "IT" + assert actual[agent.uuid]["dst_region"] == agent.region def test_aggregate_cache(): import numpy as np + import xarray as xr from pandas.testing import assert_frame_equal from muse.outputs.cache import _aggregate_cache quantity = "height" + a = xr.DataArray(np.ones((3, 4, 5)), name=quantity) b = a.copy() b[0, 0, 0] = 0 - def to_df(arr): - return arr.to_dataframe().reset_index().astype(float) - actual = _aggregate_cache(quantity, [a, b]) - assert_frame_equal(actual, to_df(b)) + assert_frame_equal(actual, b.to_dataframe().reset_index().astype(float)) actual = _aggregate_cache(quantity, [b, a]) - assert_frame_equal(actual, to_df(a)) + assert_frame_equal(actual, a.to_dataframe().reset_index().astype(float)) c = a.copy() c.assign_coords(dim_0=c.dim_0.data * 10) - dc, da = map(to_df, [c, a]) + dc, da = (da.to_dataframe().reset_index() for da in [c, a]) actual = _aggregate_cache(quantity, [c, a]) expected = pd.DataFrame.merge(dc, da, how="outer").astype(float) @@ -545,38 +516,81 @@ def to_df(arr): def test_consolidate_quantity(newcapa_agent, retro_agent): - """Test consolidation of quantity data with agent information.""" from types import SimpleNamespace from muse.outputs.cache import consolidate_quantity, extract_agents_internal - def setup_agent(agent, name, category): - agent.name = name - agent.category = category - return agent - - newcapa_agent = setup_agent(newcapa_agent, "A1", "newcapa") - retro_agent = setup_agent(retro_agent, "A2", "retro") + newcapa_agent.name = "A1" + retro_agent.name = "A2" + newcapa_agent.category = "newcapa" + retro_agent.category = "retro" sector = SimpleNamespace(name="IT", agents=[newcapa_agent, retro_agent]) agents = extract_agents_internal(sector) - def create_agent_array(agent_uuid, modify_first=False): - arr = xr.DataArray( - np.ones((3, 4, 5)), - dims=("agent", "replacement", "asset"), - coords={"agent": [agent_uuid] * 3}, - name="height", - ) - if modify_first: - arr[0, 0, 0] = 0 - return arr + quantity = "height" + a = xr.DataArray( + np.ones((3, 4, 5)), + dims=("agent", "replacement", "asset"), + coords={ + "agent": [ + newcapa_agent.uuid, + ] + * 3 + }, + name=quantity, + ) + b = a.copy() + b[0, 0, 0] = 0 + b.assign_coords( + agent=[ + retro_agent.uuid, + ] + * 3 + ) - a = create_agent_array(newcapa_agent.uuid) - b = create_agent_array(retro_agent.uuid, modify_first=True) + actual = consolidate_quantity(quantity, [a, b], agents) - actual = consolidate_quantity("height", [a, b], agents) - cols = set((*agents[retro_agent.uuid].keys(), "technology", "height")) + cols = set((*agents[retro_agent.uuid].keys(), "technology", quantity)) assert set(actual.columns) == cols assert all( name in (newcapa_agent.name, retro_agent.name) for name in actual.agent.unique() ) + + +@patch("muse.outputs.cache.consolidate_quantity") +def test_output_capacity(mock_consolidate): + import xarray as xr + + from muse.outputs.cache import capacity + + cached = [xr.DataArray() for _ in range(3)] + agents = {} + + capacity(cached, agents) + mock_consolidate.assert_called_once_with("capacity", cached, agents) + + +@patch("muse.outputs.cache.consolidate_quantity") +def test_output_production(mock_consolidate): + import xarray as xr + + from muse.outputs.cache import production + + cached = [xr.DataArray() for _ in range(3)] + agents = {} + + production(cached, agents) + mock_consolidate.assert_called_once_with("production", cached, agents) + + +@patch("muse.outputs.cache.consolidate_quantity") +def test_output_lcoe(mock_consolidate): + import xarray as xr + + from muse.outputs.cache import lcoe + + cached = [xr.DataArray() for _ in range(3)] + agents = {} + + lcoe(cached, agents) + mock_consolidate.assert_called_once_with("lcoe", cached, agents) From 25c6ba154a16ca6c77024e120995f3a964495e85 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Wed, 4 Jun 2025 10:14:39 +0100 Subject: [PATCH 24/33] Tidy test_outputs --- tests/test_outputs.py | 472 ++++++++++++++++++++---------------------- 1 file changed, 229 insertions(+), 243 deletions(-) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 10c0ab9d..89b2eb2b 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -13,6 +13,7 @@ @fixture def streetcred(save_registries): + """Create a test output quantity that returns random data.""" from muse.outputs.sector import register_output_quantity @register_output_quantity @@ -28,83 +29,103 @@ def streetcred(*args, **kwargs): ) -@mark.usefixtures("streetcred") -def test_save_with_dir(tmpdir): - from pandas import read_csv +@fixture +def market(): + """Common market fixture used in multiple tests.""" + return xr.DataArray([1], coords={"year": [2010]}, dims="year") + +@fixture +def base_config(tmpdir): + """Common config fixture used in multiple tests.""" path = Path(tmpdir) / "results" / "stuff" - config = { + return { "filename": path / "{Sector}{year}{Quantity}.csv", "quantity": "streetcred", } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here - result = factory(config, sector_name="Yoyo")(market, None, None) - assert len(result) == 1 - assert result[0] == path / "Yoyo2010Streetcred.csv" - assert result[0].exists() - assert result[0].is_file() - read_csv(result[0]) + + +def create_test_data_array(values, coords=None, name="test"): + """Helper function to create test DataArrays.""" + if coords is None: + coords = dict(a=[2, 4]) + return xr.DataArray(values, coords=coords, dims="a", name=name) + + +def assert_file_exists_and_readable(path, expected_columns=None): + """Helper to verify file exists and can be read.""" + assert path.exists() and path.is_file() + df = pd.read_csv(path) + if expected_columns: + assert set(df.columns) == set(expected_columns) + return df @mark.usefixtures("streetcred") -def test_overwrite(tmpdir): - from pytest import raises +def test_save_with_dir(tmpdir, market, base_config): + """Test saving output to directory with sector and year in filename.""" + result = factory(base_config, sector_name="Yoyo")(market, None, None) + assert len(result) == 1 + expected_path = Path(base_config["filename"]).parent / "Yoyo2010Streetcred.csv" + assert result[0] == expected_path + assert_file_exists_and_readable(result[0]) - path = Path(tmpdir) / "results" / "stuff" - config = { - "filename": path / "{Sector}{year}{Quantity}.csv", - "quantity": "streetcred", - } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here - outputter = factory(config, sector_name="Yoyo") + +@mark.usefixtures("streetcred") +def test_overwrite(tmpdir, market, base_config): + """Test file overwrite behavior.""" + outputter = factory(base_config, sector_name="Yoyo") result = outputter(market, None, None) - assert result[0] == path / "Yoyo2010Streetcred.csv" - assert result[0].is_file() + expected_path = Path(base_config["filename"]).parent / "Yoyo2010Streetcred.csv" + assert result[0] == expected_path + assert_file_exists_and_readable(result[0]) # default is to never overwrite with raises(IOError): outputter(market, None, None) - config["overwrite"] = True - factory(config, sector_name="Yoyo")(market, None, None) + base_config["overwrite"] = True + factory(base_config, sector_name="Yoyo")(market, None, None) @mark.usefixtures("streetcred") -def test_save_with_path_to_nc_with_suffix(tmpdir): +@mark.parametrize( + "config_type,suffix", + [ + ("suffix", "nc"), + ("sink", "nc"), + ], +) +def test_save_with_path_to_nc(tmpdir, market, base_config, config_type, suffix): + """Test saving output to NC file with different config types.""" path = Path(tmpdir) / "results" / "stuff" - config = { - "filename": path / "{Sector}{year}{Quantity}{suffix}", - "quantity": "streetcred", - "suffix": "nc", - } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here + if config_type == "suffix": + config = { + "filename": path / "{Sector}{year}{Quantity}{suffix}", + "quantity": "streetcred", + "suffix": suffix, + } + else: + config = { + "filename": path / "{sector}{year}{quantity}.csv", + "quantity": "streetcred", + "sink": suffix, + } result = factory(config, sector_name="Yoyo")(market, None, None) - assert result[0] == path / "Yoyo2010Streetcred.nc" - assert result[0].is_file() - xr.open_dataset(result[0]) - - -@mark.usefixtures("streetcred") -def test_save_with_path_to_nc_with_sink(tmpdir): - path = Path(tmpdir) / "results" / "stuff" - # can use None because we **know** none of the arguments are used here - config = { - "filename": path / "{sector}{year}{quantity}.csv", - "quantity": "streetcred", - "sink": "nc", - } - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - result = factory(config, sector_name="Yoyo")(market, None, None) - assert result[0] == path / "yoyo2010streetcred.csv" + expected_path = path / ( + f"{'Yoyo' if config_type == 'suffix' else 'yoyo'}" + f"2010" + f"{'Streetcred' if config_type == 'suffix' else 'streetcred'}" + f".{'nc' if config_type == 'suffix' else 'csv'}" + ) + assert result[0] == expected_path assert result[0].is_file() xr.open_dataset(result[0]) @mark.usefixtures("streetcred") -def test_save_with_fullpath_to_excel_with_sink(tmpdir): +def test_save_with_fullpath_to_excel(tmpdir, market): + """Test saving output to Excel file.""" from warnings import simplefilter from pandas import read_excel @@ -114,25 +135,38 @@ def test_save_with_fullpath_to_excel_with_sink(tmpdir): path = Path(tmpdir) / "results" / "stuff" / "this.xlsx" config = {"filename": path, "quantity": "streetcred", "sink": "xlsx"} - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - # can use None because we **know** none of the arguments are used here result = factory(config, sector_name="Yoyo")(market, None, None) assert result[0] == path - assert result[0].is_file() + assert_file_exists_and_readable(result[0]) read_excel(result[0]) -@mark.usefixtures("streetcred") -def test_no_sink_or_suffix(tmpdir): - from muse.outputs.sector import factory +@patch("muse.outputs.cache.consolidate_quantity") +def test_output_functions(mock_consolidate): + """Test output functions (capacity, production, lcoe) with common setup.""" + from muse.outputs.cache import capacity, lcoe, production + + cached = [xr.DataArray() for _ in range(3)] + agents = {} + + for func, quantity in [ + (capacity, "capacity"), + (production, "production"), + (lcoe, "lcoe"), + ]: + func(cached, agents) + mock_consolidate.assert_called_once_with(quantity, cached, agents) + mock_consolidate.reset_mock() + +@mark.usefixtures("streetcred") +def test_no_sink_or_suffix(tmpdir, market): + """Test default sink and suffix behavior.""" config = dict( quantity="streetcred", filename=f"{tmpdir}/{{Sector}}{{Quantity}}{{year}}{{suffix}}", ) - outputs = factory(config) - market = xr.DataArray([1], coords={"year": [2010]}, dims="year") - result = outputs(market, None, None) + result = factory(config)(market, None, None) assert len(result) == 1 assert result[0].is_file() assert result[0].suffix == ".csv" @@ -140,6 +174,7 @@ def test_no_sink_or_suffix(tmpdir): @mark.usefixtures("save_registries") def test_can_register_class(): + """Test class registration functionality.""" from muse.outputs.sinks import factory, register_output_sink @register_output_sink @@ -151,12 +186,14 @@ def __init__(self, sector, some_args=3): def __call__(self, x): pass + # Test default arguments settings = {"sink": {"name": "AClass"}} sink = factory(settings, sector_name="yoyo") assert isinstance(sink, AClass) assert sink.sector == "yoyo" assert sink.some_args == 3 + # Test custom arguments settings = {"sink": {"name": "AClass", "some_args": 5}} sink = factory(settings, sector_name="yoyo") assert isinstance(sink, AClass) @@ -166,6 +203,7 @@ def __call__(self, x): @mark.usefixtures("save_registries") def test_can_register_function(): + """Test function registration functionality.""" from muse.outputs.sinks import factory, register_output_sink @register_output_sink @@ -179,8 +217,10 @@ def a_function(*args): @mark.usefixtures("save_registries") def test_yearly_aggregate(): + """Test yearly aggregation with custom sink.""" from muse.outputs.sinks import factory, register_output_sink + # Setup tracking variables received_data = None gyear = None gsector = None @@ -202,16 +242,17 @@ def dummy(data, year: int, sector: str, overwrite: bool) -> MySpecialReturn: dict(overwrite=True, sink=dict(aggregate="dummy")), sector_name="yoyo" ) - data = xr.DataArray([1, 0], coords=dict(a=[2, 4]), dims="a", name="nada") + # Test first year + data = create_test_data_array([1, 0], name="nada") data["year"] = 2010 - assert isinstance(sink(data, 2010), MySpecialReturn) assert gyear == 2010 assert gsector == "yoyo" assert goverwrite is True assert isinstance(received_data, pd.DataFrame) - data = xr.DataArray([0, 1], coords=dict(a=[2, 4]), dims="a", name="nada") + # Test second year + data = create_test_data_array([0, 1], name="nada") data["year"] = 2020 assert isinstance(sink(data, 2020), MySpecialReturn) assert gyear == 2020 @@ -221,34 +262,42 @@ def dummy(data, year: int, sector: str, overwrite: bool) -> MySpecialReturn: def test_yearly_aggregate_file(tmpdir): + """Test yearly aggregation to file with multiple years of data.""" from muse.outputs.sinks import factory path = Path(tmpdir) / "file.csv" sink = factory(dict(filename=str(path), sink="aggregate"), sector_name="yoyo") - data = xr.DataArray([1, 0], coords=dict(a=[2, 4]), dims="a", name="georges") - data["year"] = 2010 - assert sink(data, 2010) == path - dataframe = pd.read_csv(path) - assert set(dataframe.columns) == {"year", "georges"} - assert dataframe.shape[0] == 2 + def verify_year_data(values, year, expected_rows): + data = create_test_data_array(values, name="georges") + data["year"] = year + assert sink(data, year) == path + df = assert_file_exists_and_readable(path, {"year", "georges"}) + assert df.shape[0] == expected_rows + return df - data = xr.DataArray([0, 1], coords=dict(a=[2, 4]), dims="a", name="georges") - data["year"] = 2020 - assert sink(data, 2020) == path - dataframe = pd.read_csv(path) - assert set(dataframe.columns) == {"year", "georges"} - assert dataframe.shape[0] == 4 + # Test first year + verify_year_data([1, 0], 2010, 2) + + # Test second year (should append to existing file) + df2 = verify_year_data([0, 1], 2020, 4) + + # Verify data from both years is present + assert set(df2.year.unique()) == {2010, 2020} + assert df2[df2.year == 2010].georges.tolist() == [1, 0] + assert df2[df2.year == 2020].georges.tolist() == [0, 1] def test_yearly_aggregate_no_outputs(tmpdir): + """Test behavior with no outputs configured.""" from muse.outputs.mca import factory outputs = factory() assert len(outputs(None, year=2010)) == 0 -def test_mca_aggregate_outputs(tmpdir): +def setup_mca_test(tmpdir, outputs_config): + """Helper function to set up MCA tests.""" from toml import dump, load from muse import examples @@ -256,17 +305,20 @@ def test_mca_aggregate_outputs(tmpdir): examples.copy_model(path=str(tmpdir)) settings = load(str(tmpdir / "model" / "settings.toml")) - settings["outputs"] = [ - dict(filename="{path}/{Quantity}{suffix}", quantity="prices", sink="aggregate") - ] + settings["outputs"] = [outputs_config] settings["time_framework"] = settings["time_framework"][:2] dump(settings, (tmpdir / "model" / "settings.toml")) + return MCA.factory(str(tmpdir / "model" / "settings.toml")) - mca = MCA.factory(str(tmpdir / "model" / "settings.toml")) - mca.run() +def test_mca_aggregate_outputs(tmpdir): + """Test MCA aggregate outputs.""" + mca = setup_mca_test( + tmpdir, + dict(filename="{path}/{Quantity}{suffix}", quantity="prices", sink="aggregate"), + ) + mca.run() assert (tmpdir / "model" / "Prices.csv").exists() - # TODO: should pass again after #612 # data = pd.read_csv(tmpdir / "model" / "Prices.csv") # assert set(data.year) == set(settings["time_framework"]) @@ -274,23 +326,10 @@ def test_mca_aggregate_outputs(tmpdir): @mark.usefixtures("save_registries") def test_path_formatting(tmpdir): - from toml import dump, load - - from muse.examples import copy_model - from muse.mca import MCA + """Test path formatting with dummy sink and quantity.""" from muse.outputs.mca import register_output_quantity from muse.outputs.sinks import register_output_sink, sink_to_file - # Copy the data to tmpdir - copy_model(path=tmpdir) - - settings_file = tmpdir / "model" / "settings.toml" - settings = load(settings_file) - settings["outputs"] = [ - dict(quantity="dummy", sink="to_dummy", filename="{path}/{Quantity}{suffix}") - ] - dump(settings, (settings_file)) - @register_output_sink(name="dummy_sink") @sink_to_file(".dummy") def to_dummy(quantity, filename, **params) -> None: @@ -300,12 +339,11 @@ def to_dummy(quantity, filename, **params) -> None: def dummy(market, **kwargs): return xr.DataArray() - mca = MCA.factory(Path(settings_file)) - assert mca.outputs(mca.market)[0] == Path( - settings["outputs"][0]["filename"].format( - path=tmpdir / "model", Quantity="Dummy", suffix=".dummy" - ) + mca = setup_mca_test( + tmpdir, + dict(quantity="dummy", sink="to_dummy", filename="{path}/{Quantity}{suffix}"), ) + assert mca.outputs(mca.market)[0] == Path(tmpdir / "model" / "Dummy.dummy") def test_register_output_quantity_cache(): @@ -319,64 +357,59 @@ def dummy_quantity(*args): class TestOutputCache: + @fixture + def output_params(self): + return [dict(quantity="height"), dict(quantity="width")] + + @fixture + def output_quantities(self, output_params): + quantities = {q["quantity"]: lambda _: None for q in output_params} + quantities["depth"] = lambda _: None + return quantities + + @fixture + def topic(self): + return "BBC Muse" + @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_init(self, mock_factory, mock_subscribe): + def test_init( + self, mock_factory, mock_subscribe, output_params, output_quantities, topic + ): from muse.outputs.cache import OutputCache - param = [dict(quantity="height"), dict(quantity="width")] - output_quantities = {q["quantity"]: lambda _: None for q in param} - output_quantities["depth"] = lambda _: None - topic = "BBC Muse" - output_cache = OutputCache( - *param, output_quantities=output_quantities, topic=topic + *output_params, output_quantities=output_quantities, topic=topic ) - - assert mock_factory.call_count == len(param) + assert mock_factory.call_count == len(output_params) mock_subscribe.assert_called_once_with(output_cache.cache, topic) @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_cache(self, mock_factory, mock_subscribe): - import xarray as xr - + def test_cache( + self, mock_factory, mock_subscribe, output_params, output_quantities, topic + ): from muse.outputs.cache import OutputCache - param = [dict(quantity="height"), dict(quantity="width")] - output_quantities = {q["quantity"]: lambda _: None for q in param} - output_quantities["depth"] = lambda _: None - topic = "BBC Muse" - output_cache = OutputCache( - *param, output_quantities=output_quantities, topic=topic + *output_params, output_quantities=output_quantities, topic=topic ) - output_cache.cache(dict(height=xr.DataArray(), depth=xr.DataArray())) - assert len(output_cache.to_save.get("height")) == 1 assert len(output_cache.to_save.get("depth", [])) == 0 @patch("pubsub.pub.subscribe") @patch("muse.outputs.sector._factory") - def test_consolidate_cache(self, mock_factory, mock_subscribe): - import xarray as xr - + def test_consolidate_cache( + self, mock_factory, mock_subscribe, output_params, output_quantities, topic + ): from muse.outputs.cache import OutputCache - param = [dict(quantity="height"), dict(quantity="width")] - output_quantities = {q["quantity"]: lambda _: None for q in param} - output_quantities["depth"] = lambda _: None - topic = "BBC Muse" - year = 2042 - output_cache = OutputCache( - *param, output_quantities=output_quantities, topic=topic + *output_params, output_quantities=output_quantities, topic=topic ) - output_cache.cache(dict(height=xr.DataArray())) - output_cache.consolidate_cache(year) - + output_cache.consolidate_cache(2042) output_cache.factory["height"].assert_called_once() @@ -388,9 +421,12 @@ def test_cache_quantity(mock_match, mock_send): result = {"mass": 42} mock_match.return_value = result + def verify_message_sent(): + mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) + mock_send.reset_mock() + cache_quantity(**result) - mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) - mock_send.reset_mock() + verify_message_sent() with raises(ValueError): cache_quantity(function=lambda: None, mass=42) @@ -407,7 +443,7 @@ def fun2(): fun2() mock_match.assert_called_once_with("mass", 42) - mock_send.assert_called_once_with(CACHE_TOPIC_CHANNEL, data=result) + verify_message_sent() def test_match_quantities(): @@ -415,37 +451,27 @@ def test_match_quantities(): from muse.outputs.cache import match_quantities - q = "mass" - da = xr.DataArray(name=q) - ds = xr.Dataset({q: da}) - def assert_equal(a: dict[str, xr.DataArray], b: dict[str, xr.DataArray]): assert set(a.keys()) == set(b.keys()) for k in a: xr.testing.assert_equal(a[k], b[k]) - actual = match_quantities(quantity=q, data=da) - assert_equal(actual, {q: da}) - - actual = match_quantities(quantity=q, data=ds) - assert_equal(actual, {q: da}) + # Test single quantity with DataArray + q = "mass" + da = xr.DataArray(name=q) + ds = xr.Dataset({q: da}) + assert_equal(match_quantities(quantity=q, data=da), {q: da}) + assert_equal(match_quantities(quantity=q, data=ds), {q: da}) + # Test multiple quantities with Dataset p = "height" ds = xr.Dataset({q: da, p: da, "rubish": da}) - actual = match_quantities(quantity=[q, p], data=ds) - assert_equal(actual, {q: da, p: da}) - - actual = match_quantities(quantity=[q, p], data=[da, da]) - assert_equal(actual, {q: da, p: da}) + assert_equal(match_quantities(quantity=[q, p], data=ds), {q: da, p: da}) + assert_equal(match_quantities(quantity=[q, p], data=[da, da]), {q: da, p: da}) + # Test error cases with raises(ValueError): - match_quantities( - quantity=[q, p], - data=[ - da, - ], - ) - + match_quantities(quantity=[q, p], data=[da]) with raises(TypeError): match_quantities(quantity=[q, p], data=42) @@ -464,51 +490,54 @@ def test_extract_agents(mock_extract): def test_extract_agents_internal(newcapa_agent, retro_agent): + """Test internal agent extraction.""" from types import SimpleNamespace from muse.outputs.cache import extract_agents_internal - newcapa_agent.name = "A1" - retro_agent.name = "A2" - sector = SimpleNamespace(name="IT", agents=[newcapa_agent, retro_agent]) + def setup_agent(agent, name): + agent.name = name + return agent + + agents = [setup_agent(newcapa_agent, "A1"), setup_agent(retro_agent, "A2")] + sector = SimpleNamespace(name="IT", agents=agents) actual = extract_agents_internal(sector) - for agent in [newcapa_agent, retro_agent]: + expected_keys = ("agent", "category", "sector", "dst_region") + + for agent in agents: assert agent.uuid in actual - assert tuple(actual[agent.uuid].keys()) == ( - "agent", - "category", - "sector", - "dst_region", - ) - assert actual[agent.uuid]["agent"] == agent.name - assert actual[agent.uuid]["category"] == agent.category - assert actual[agent.uuid]["sector"] == "IT" - assert actual[agent.uuid]["dst_region"] == agent.region + assert tuple(actual[agent.uuid].keys()) == expected_keys + agent_data = actual[agent.uuid] + assert agent_data["agent"] == agent.name + assert agent_data["category"] == agent.category + assert agent_data["sector"] == "IT" + assert agent_data["dst_region"] == agent.region def test_aggregate_cache(): import numpy as np - import xarray as xr from pandas.testing import assert_frame_equal from muse.outputs.cache import _aggregate_cache quantity = "height" - a = xr.DataArray(np.ones((3, 4, 5)), name=quantity) b = a.copy() b[0, 0, 0] = 0 + def to_df(arr): + return arr.to_dataframe().reset_index().astype(float) + actual = _aggregate_cache(quantity, [a, b]) - assert_frame_equal(actual, b.to_dataframe().reset_index().astype(float)) + assert_frame_equal(actual, to_df(b)) actual = _aggregate_cache(quantity, [b, a]) - assert_frame_equal(actual, a.to_dataframe().reset_index().astype(float)) + assert_frame_equal(actual, to_df(a)) c = a.copy() c.assign_coords(dim_0=c.dim_0.data * 10) - dc, da = (da.to_dataframe().reset_index() for da in [c, a]) + dc, da = map(to_df, [c, a]) actual = _aggregate_cache(quantity, [c, a]) expected = pd.DataFrame.merge(dc, da, how="outer").astype(float) @@ -516,81 +545,38 @@ def test_aggregate_cache(): def test_consolidate_quantity(newcapa_agent, retro_agent): + """Test consolidation of quantity data with agent information.""" from types import SimpleNamespace from muse.outputs.cache import consolidate_quantity, extract_agents_internal - newcapa_agent.name = "A1" - retro_agent.name = "A2" - newcapa_agent.category = "newcapa" - retro_agent.category = "retro" + def setup_agent(agent, name, category): + agent.name = name + agent.category = category + return agent + + newcapa_agent = setup_agent(newcapa_agent, "A1", "newcapa") + retro_agent = setup_agent(retro_agent, "A2", "retro") sector = SimpleNamespace(name="IT", agents=[newcapa_agent, retro_agent]) agents = extract_agents_internal(sector) - quantity = "height" - a = xr.DataArray( - np.ones((3, 4, 5)), - dims=("agent", "replacement", "asset"), - coords={ - "agent": [ - newcapa_agent.uuid, - ] - * 3 - }, - name=quantity, - ) - b = a.copy() - b[0, 0, 0] = 0 - b.assign_coords( - agent=[ - retro_agent.uuid, - ] - * 3 - ) + def create_agent_array(agent_uuid, modify_first=False): + arr = xr.DataArray( + np.ones((3, 4, 5)), + dims=("agent", "replacement", "asset"), + coords={"agent": [agent_uuid] * 3}, + name="height", + ) + if modify_first: + arr[0, 0, 0] = 0 + return arr - actual = consolidate_quantity(quantity, [a, b], agents) + a = create_agent_array(newcapa_agent.uuid) + b = create_agent_array(retro_agent.uuid, modify_first=True) - cols = set((*agents[retro_agent.uuid].keys(), "technology", quantity)) + actual = consolidate_quantity("height", [a, b], agents) + cols = set((*agents[retro_agent.uuid].keys(), "technology", "height")) assert set(actual.columns) == cols assert all( name in (newcapa_agent.name, retro_agent.name) for name in actual.agent.unique() ) - - -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_capacity(mock_consolidate): - import xarray as xr - - from muse.outputs.cache import capacity - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - capacity(cached, agents) - mock_consolidate.assert_called_once_with("capacity", cached, agents) - - -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_production(mock_consolidate): - import xarray as xr - - from muse.outputs.cache import production - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - production(cached, agents) - mock_consolidate.assert_called_once_with("production", cached, agents) - - -@patch("muse.outputs.cache.consolidate_quantity") -def test_output_lcoe(mock_consolidate): - import xarray as xr - - from muse.outputs.cache import lcoe - - cached = [xr.DataArray() for _ in range(3)] - agents = {} - - lcoe(cached, agents) - mock_consolidate.assert_called_once_with("lcoe", cached, agents) From cdc0d8309daecc4d5db29293a72605edb4e110fa Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Thu, 5 Jun 2025 18:12:45 +0100 Subject: [PATCH 25/33] Fix test_save_with_fullpath_to_excel --- tests/test_outputs.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 89b2eb2b..6becd263 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -124,8 +124,7 @@ def test_save_with_path_to_nc(tmpdir, market, base_config, config_type, suffix): @mark.usefixtures("streetcred") -def test_save_with_fullpath_to_excel(tmpdir, market): - """Test saving output to Excel file.""" +def test_save_with_fullpath_to_excel(tmpdir): from warnings import simplefilter from pandas import read_excel @@ -137,7 +136,7 @@ def test_save_with_fullpath_to_excel(tmpdir, market): config = {"filename": path, "quantity": "streetcred", "sink": "xlsx"} result = factory(config, sector_name="Yoyo")(market, None, None) assert result[0] == path - assert_file_exists_and_readable(result[0]) + assert result[0].is_file() read_excel(result[0]) From 63eae992df77c3543752986eddb2928cd4d3b5b5 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Thu, 5 Jun 2025 18:19:29 +0100 Subject: [PATCH 26/33] Actually fix test --- tests/test_outputs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 6becd263..ed546531 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -134,6 +134,8 @@ def test_save_with_fullpath_to_excel(tmpdir): path = Path(tmpdir) / "results" / "stuff" / "this.xlsx" config = {"filename": path, "quantity": "streetcred", "sink": "xlsx"} + market = xr.DataArray([1], coords={"year": [2010]}, dims="year") + # can use None because we **know** none of the arguments are used here result = factory(config, sector_name="Yoyo")(market, None, None) assert result[0] == path assert result[0].is_file() From 1cfc24fae7024d9ff96a0d443f206d3eb4f8deba Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Mon, 2 Jun 2025 23:01:33 +0100 Subject: [PATCH 27/33] Tidy lp adapter, move to new module from __future__ import annotations Restore some deleted test behaviour Split tests Revert tests --- src/muse/constraints.py | 479 +--------------------------------------- src/muse/investments.py | 4 +- src/muse/lp_adapter.py | 354 +++++++++++++++++++++++++++++ 3 files changed, 357 insertions(+), 480 deletions(-) create mode 100644 src/muse/lp_adapter.py diff --git a/src/muse/constraints.py b/src/muse/constraints.py index be0dcbc3..5c2aa687 100644 --- a/src/muse/constraints.py +++ b/src/muse/constraints.py @@ -94,7 +94,6 @@ def constraints( from __future__ import annotations from collections.abc import Mapping, MutableMapping, Sequence -from dataclasses import dataclass from enum import Enum, auto from typing import ( Any, @@ -110,7 +109,7 @@ def constraints( from mypy_extensions import KwArg from muse.registration import registrator -from muse.timeslices import broadcast_timeslice, distribute_timeslice, drop_timeslice +from muse.timeslices import broadcast_timeslice, distribute_timeslice CAPACITY_DIMS = "asset", "replacement", "region" """Default dimensions for capacity decision variables.""" @@ -752,479 +751,3 @@ def minimum_service( dict(capacity=-capacity, production=production, b=b), attrs=dict(kind=ConstraintKind.LOWER_BOUND), ) - - -def lp_costs( - capacity_costs: xr.DataArray, - commodities: list[str], - timeslice_level: str | None = None, -) -> xr.Dataset: - """Creates dataset of costs for solving with scipy's LP solver. - - Importantly, this also defines the decision variables in the linear program. - - The costs applied to the capacity decision variables are provided. This should - have dimensions "asset" and "replacement". In other words, capacity addition - is solved for each replacement technology for each existing asset. - - No cost is applied to the production decision variables. Thus, the production - component of the resulting dataset is zero, with dimensions determining the - production decision variables. This will have dimensions "asset", "replacement", - "commodity", and "timeslice". In other words, production is solved for each - replacement technology for each existing asset, for each commodity, and for each - timeslice. - - Args: - capacity_costs: DataArray with dimensions "asset" and "replacement" defining the - costs of adding capacity to the system. - commodities: List of commodities to create production decision variables for. - timeslice_level: The timeslice level of the linear problem. - """ - assert set(capacity_costs.dims) == {"asset", "replacement"} - - # Start with capacity costs as template (defines "asset" and "replacement" dims) - production_costs = xr.zeros_like(capacity_costs) - - # Add a "timeslice" dimension, convert multiindex to single index - production_costs = broadcast_timeslice(production_costs, level=timeslice_level) - production_costs = drop_timeslice(production_costs) - production_costs["timeslice"] = pd.Index( - production_costs.get_index("timeslice"), tupleize_cols=False - ) - - # Add a "commodity" dimension - production_costs = production_costs.expand_dims(commodity=commodities) - assert set(production_costs.dims) == { - "asset", - "replacement", - "commodity", - "timeslice", - } - - # Result is dataset of provided capacity costs and zero production costs - return xr.Dataset(dict(capacity=capacity_costs, production=production_costs)) - - -def lp_constraint(constraint: Constraint, lpcosts: xr.Dataset) -> Constraint: - """Transforms the constraint to LP data. - - The goal is to create from ``lpcosts.capacity``, ``constraint.capacity``, and - ``constraint.b`` a 2d-matrix ``constraint`` vs ``decision variables``. - - #. The dimensions of ``constraint.b`` are the constraint dimensions. They are - renamed ``"c(xxx)"``. - #. The dimensions of ``lpcosts`` are the decision-variable dimensions. They are - renamed ``"d(xxx)"``. - #. ``set(b.dims).intersection(lpcosts.xxx.dims)`` are diagonal - in constraint dimensions and decision variables dimension, with ``xxx`` the - capacity or the production - #. ``set(constraint.xxx.dims) - set(lpcosts.xxx.dims) - set(b.dims)`` are reduced by - summation, with ``xxx`` the capacity or the production - #. ``set(lpcosts.xxx.dims) - set(constraint.xxx.dims) - set(b.dims)`` are added for - expansion, with ``xxx`` the capacity or the production - - See :py:func:`muse.constraints.lp_constraint_matrix` for a more detailed explanation - of the transformations applied here. - """ - constraint = constraint.copy(deep=False) - - # Deal with timeslice multiindex - if "timeslice" in constraint.dims: - constraint = drop_timeslice(constraint) - constraint["timeslice"] = pd.Index( - constraint.get_index("timeslice"), tupleize_cols=False - ) - - # Rename dimensions in b - b = constraint.b.drop_vars(set(constraint.b.coords) - set(constraint.b.dims)) - b = b.rename({k: f"c({k})" for k in b.dims}) - - # Create capacity constraint matrix - capacity = lp_constraint_matrix(constraint.b, constraint.capacity, lpcosts.capacity) - capacity = capacity.drop_vars(set(capacity.coords) - set(capacity.dims)) - - # Create production constraint matrix - production = lp_constraint_matrix( - constraint.b, constraint.production, lpcosts.production - ) - production = production.drop_vars(set(production.coords) - set(production.dims)) - - # Combine data - result = xr.Dataset( - {"b": b, "capacity": capacity, "production": production}, attrs=constraint.attrs - ) - return result - - -def lp_constraint_matrix( - b: xr.DataArray, constraint: xr.DataArray, lpcosts: xr.DataArray -): - """Transforms one constraint block into an lp matrix. - - The goal is to create from ``lpcosts``, ``constraint``, and ``b`` a 2d-matrix of - constraints vs decision variables. - - #. The dimensions of ``b`` are the constraint dimensions. They are renamed - ``"c(xxx)"``. - #. The dimensions of ``lpcosts`` are the decision-variable dimensions. They are - renamed ``"d(xxx)"``. - #. ``set(b.dims).intersection(lpcosts.dims)`` are diagonal - in constraint dimensions and decision variables dimension - #. ``set(constraint.dims) - set(lpcosts.dims) - set(b.dims)`` are reduced by - summation - #. ``set(lpcosts.dims) - set(constraint.dims) - set(b.dims)`` are added for - expansion - #. ``set(b.dims) - set(constraint.dims) - set(lpcosts.dims)`` are added for - expansion. Such dimensions only make sense if they consist of one point. - - The result is the constraint matrix, expanded, reduced and diagonalized for the - conditions above. - - """ - from functools import reduce - - from numpy import eye - - # Sum over all dimensions that are not in the constraint or the decision variables - result = constraint.sum(set(constraint.dims) - set(lpcosts.dims) - set(b.dims)) - - # Rename dimensions for decision variables - result = result.rename( - {k: f"d({k})" for k in set(result.dims).intersection(lpcosts.dims)} - ) - - # Rename dimensions for constraints - result = result.rename( - {k: f"c({k})" for k in set(result.dims).intersection(b.dims)} - ) - - # Expand dimensions that are in the decision variables but not in the constraint - expand = set(lpcosts.dims) - set(constraint.dims) - set(b.dims) - result = result.expand_dims( - {f"d({k})": lpcosts[k].rename({k: f"d({k})"}).set_index() for k in expand} - ) - - # Expand dimensions that are in the constraint but not in the decision variables - expand = set(b.dims) - set(constraint.dims) - set(lpcosts.dims) - result = result.expand_dims( - {f"c({k})": b[k].rename({k: f"c({k})"}).set_index() for k in expand} - ) - - # Dimensions that are in both the decision variables and the constraint - diag_dims = set(b.dims).intersection(lpcosts.dims) - diag_dims = sorted(diag_dims) - - if diag_dims: - - def get_dimension(dim): - if dim in b.dims: - return b[dim].values - if dim in lpcosts.dims: - return lpcosts[dim].values - return constraint[dim].values - - diagonal_submats = [ - xr.DataArray( - eye(len(b[k])), - coords={f"c({k})": get_dimension(k), f"d({k})": get_dimension(k)}, - dims=(f"c({k})", f"d({k})"), - ) - for k in diag_dims - ] - reduced = reduce(xr.DataArray.__mul__, diagonal_submats) - if "d(timeslice)" in reduced.dims: - reduced = reduced.drop_vars("d(timeslice)") - result = result * reduced - - return result - - -@dataclass -class ScipyAdapter: - """Creates the input for the scipy solvers. - - This adapter is required to convert data (costs and constraints) from a series of - dataarrays to a format that can be used by scipy's linear programming solver. - - The scipy solver requires the following inputs as a set of 1D or 2D numpy arrays: - c: Coefficients of the linear objective function to be minimized. This is a 1D - vector, each element representing the cost of a decision variable. - A_ub: The inequality constraint matrix. This is a 2D matrix containing - coefficients of the linear inequality constraints, with constraints as rows - and decision variables as columns. - b_ub: The inequality constraint vector. This is a 1D vector, each element - representing an upper bound on the corresponding constraint in A_ub. - A_eq: The equality constraint matrix. In practice, since all of the constraints - currently implemented are inequality constraints, this will always be zeros. - However, this may be changed in the future. - b_eq: The equality constraint vector. As above, this will currently always be - zeros. - - See the documentation for `scipy.optimize.linprog` for more details: - https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.linprog.html - - This class takes in costs and constraints as xarray dataarrays, and performs the - necessary conversions to create the above inputs. These are saved together in the - `kwargs` property, which is passed to the scipy solver. This has an additional - attribute `bounds`, which is a tuple of the lower and upper bounds for the - decision variables. In practice, we will always use (0, np.inf) as the bounds. - - The output of the scipy solver is a 1D array of decision variables, which must then - be converted back into a format that MUSE can understand. To aid this, we - save templates for the capacity and production decision variables, which are used - to convert the output back into a labelled xarray dataset using the helper function - `to_muse`. - """ - - c: np.ndarray - to_muse: Callable[[np.ndarray], xr.Dataset] - bounds: tuple[float | None, float | None] = (0, np.inf) - A_ub: np.ndarray | None = None - b_ub: np.ndarray | None = None - A_eq: np.ndarray | None = None - b_eq: np.ndarray | None = None - - @classmethod - def factory( - cls, - costs: xr.DataArray, - constraints: list[Constraint], - commodities: list[str], - timeslice_level: str | None = None, - ) -> ScipyAdapter: - # Calculate costs for the linear problem - lpcosts = lp_costs( - capacity_costs=costs, - commodities=commodities, - timeslice_level=timeslice_level, - ) - - # Create dataset from costs and constraints - data = cls._unified_dataset(lpcosts, *constraints) - - # Get capacity constraint matrix / costs - capacities = cls._selected_quantity(data, "capacity") - - # Get production constraint matrix / costs - productions = cls._selected_quantity(data, "production") - - # Get constraint vector - bs = cls._selected_quantity(data, "b") - - # Prepare scipy adapter from constraints - kwargs = cls._to_scipy_adapter(capacities, productions, bs, *constraints) - - def to_muse(x: np.ndarray) -> xr.Dataset: - return ScipyAdapter._back_to_muse( - x, - capacity_template=capacities.costs, - production_template=productions.costs, - ) - - return ScipyAdapter(to_muse=to_muse, **kwargs) - - @property - def kwargs(self): - return { - "c": self.c, - "A_eq": self.A_eq, - "b_eq": self.b_eq, - "A_ub": self.A_ub, - "b_ub": self.b_ub, - "bounds": self.bounds, - } - - @staticmethod - def _unified_dataset(lpcosts: xr.Dataset, *constraints: Constraint) -> xr.Dataset: - """Creates single xr.Dataset from costs and constraints.""" - from xarray import merge - - # Reformat constraints to lp format - lp_constraints = [ - lp_constraint(constraint, lpcosts) for constraint in constraints - ] - - # Rename variables in lp constraints - lp_constraints = [ - constraint.rename( - b=f"b{i}", capacity=f"capacity{i}", production=f"production{i}" - ) - for i, constraint in enumerate(lp_constraints) - ] - - # Rename dimensions in lpcosts - lpcosts = lpcosts.rename({k: f"d({k})" for k in lpcosts.dims}) - - # Merge data - data = merge([lpcosts, *lp_constraints]) - - # An adjustment is required for lower bound constraints - for i, constraint in enumerate(constraints): - if constraint.kind == ConstraintKind.LOWER_BOUND: - data[f"b{i}"] = -data[f"b{i}"] - data[f"capacity{i}"] = -data[f"capacity{i}"] - data[f"production{i}"] = -data[f"production{i}"] - - # Enusure consistent ordering of dimensions - return data.transpose(*data.dims) - - @staticmethod - def _selected_quantity(data: xr.Dataset, name: str) -> xr.Dataset: - # Select data for the specified quantity ("capacity", "production", or "b") - result = cast( - xr.Dataset, data[[u for u in data.data_vars if str(u).startswith(name)]] - ) - - # Rename variables ("costs" for the costs variable, 0/1/2 etc. for constraints) - return result.rename( - { - k: ("costs" if k == name else int(str(k).replace(name, ""))) - for k in result.data_vars - } - ) - - @staticmethod - def _to_scipy_adapter( - capacities: xr.Dataset, productions: xr.Dataset, bs: xr.Dataset, *constraints - ): - """Converts constraints to scipy format. - - The constraints are converted to a format that can be used by scipy's linear - programming solver. The constraints are converted to a 2D matrix of constraints - vs decision variables. The constraints are then converted to a dictionary that - can be used by scipy's linear programming solver. - - Args: - capacities: Dataset with decision variables for capacity constraints. - productions: Dataset with decision variables for production constraints. - bs: Dataset with constraints. - *constraints: List of constraints. - - Returns: - Dictionary with constraints in a format that can be used by scipy's linear - programming solver. - """ - - def reshape(matrix: xr.DataArray) -> np.ndarray: - """Convert constraints matrix to a 2D np array. - - The rows of the constraaints matrix will represent the constraints, and the - columns will represent the decision variables. - """ - # Before building LP we need to sort dimensions for consistency - if list(matrix.dims) != sorted(matrix.dims): - matrix = matrix.transpose(*sorted(matrix.dims)) - - # Size of the first dimension - # This dimension represents the number of constraints - size = np.prod( - [matrix[u].shape[0] for u in matrix.dims if str(u).startswith("c")] - ) - - # Reshape into a 2D array: N constraints x N decision variables - return matrix.values.reshape((size, -1)) - - def extract_bA(constraints, *kinds: ConstraintKind): - """Extracts A and b for constraints of specified kinds. - - These will end up as A_ub and b_ub for inequality constraints, and - A_eq and b_eq for equality constraints (see ScipyAdapter). - """ - # Get indices of constraints of the specified kind - indices = [i for i in range(len(bs)) if constraints[i].kind in kinds] - - # Convert constraints matrices to 2d np arrays - capa_constraints = [reshape(capacities[i]) for i in indices] - prod_constraints = [reshape(productions[i]) for i in indices] - - # Convert constraints vectors to 1d - constraints_vectors = [ - bs[i].stack(constraint=sorted(bs[i].dims)) for i in indices - ] - - # Concatenate constraints - if capa_constraints: - A = np.concatenate( - ( - np.concatenate(capa_constraints, axis=0), - np.concatenate(prod_constraints, axis=0), - ), - axis=1, - ) - b = np.concatenate(constraints_vectors, axis=0) - else: - # If there are no constraints of the given kind, return None - A = None - b = None - return A, b - - # Create costs vector by concatenating capacity and production costs - c = np.concatenate( - ( - capacities["costs"].values.flatten(), - productions["costs"].values.flatten(), - ), - axis=0, - ) - - # Extract A and b for inequality constraints - A_ub, b_ub = extract_bA( - constraints, ConstraintKind.UPPER_BOUND, ConstraintKind.LOWER_BOUND - ) - - # Extract A and b for equality constraints - A_eq, b_eq = extract_bA(constraints, ConstraintKind.EQUALITY) - - # Prepare scipy adapter - return { - "c": c, - "A_ub": A_ub, - "b_ub": b_ub, - "A_eq": A_eq, - "b_eq": b_eq, - "bounds": (0, np.inf), - } - - @staticmethod - def _back_to_muse_quantity( - x: np.ndarray, template: xr.DataArray | xr.Dataset - ) -> xr.DataArray: - """Convert a vector of decision variables to a DataArray. - - Args: - x: 1D vector of decision variables, outputted from the scipy solver. - template: Template for the decision variables. This may be for either - capacity or production variables. - """ - # First create a multidimensional dataarray based on the template - result = xr.DataArray( - x.reshape(template.shape), coords=template.coords, dims=template.dims - ) - - # Then rename the dimensions (e.g. "d(asset)" -> "asset") - return result.rename({k: str(k)[2:-1] for k in result.dims}) - - @staticmethod - def _back_to_muse( - x: np.ndarray, - capacity_template: xr.DataArray, - production_template: xr.DataArray, - ) -> xr.Dataset: - """Convert the full set of decision variables to a Dataset. - - This must have capacity variables first, followed by production variables. - - Args: - x: 1D vector of decision variables, outputted from the scipy solver. - capacity_template: Template for the capacity decision variables. - production_template: Template for the production decision variables. - """ - n_capa = capacity_template.size # number of capacity decision variables - - capa = ScipyAdapter._back_to_muse_quantity( - x[:n_capa], template=capacity_template - ) - prod = ScipyAdapter._back_to_muse_quantity( - x[n_capa:], template=production_template - ) - return xr.Dataset({"capacity": capa, "production": prod}) diff --git a/src/muse/investments.py b/src/muse/investments.py index f5467d91..78be4490 100644 --- a/src/muse/investments.py +++ b/src/muse/investments.py @@ -288,7 +288,7 @@ def scipy_match_demand( from scipy.optimize import linprog - from muse.constraints import ScipyAdapter + from muse.lp_adapter import ScipyAdapter assert "year" not in technologies.dims @@ -296,7 +296,7 @@ def scipy_match_demand( costs = timeslice_max(costs) # Run scipy optimization with highs solver - adapter = ScipyAdapter.factory( + adapter = ScipyAdapter.from_muse_data( costs=costs, constraints=constraints, commodities=commodities, diff --git a/src/muse/lp_adapter.py b/src/muse/lp_adapter.py new file mode 100644 index 00000000..27e5615a --- /dev/null +++ b/src/muse/lp_adapter.py @@ -0,0 +1,354 @@ +"""Utilities for adapting MUSE data structures to linear programming solvers. + +This module provides utilities to convert MUSE's xarray-based data structures to and +from the format required by scipy's linear programming solver. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pandas as pd +import xarray as xr + +from muse.constraints import ConstraintKind +from muse.timeslices import broadcast_timeslice, drop_timeslice + + +def unified_dataset(lpcosts: xr.Dataset, *constraints) -> xr.Dataset: + """Creates single xr.Dataset from costs and constraints.""" + from xarray import merge + + # Reformat constraints to lp format + lp_constraints = [lp_constraint(constraint, lpcosts) for constraint in constraints] + + # Rename variables in lp constraints + lp_constraints = [ + constraint.rename( + b=f"b{i}", capacity=f"capacity{i}", production=f"production{i}" + ) + for i, constraint in enumerate(lp_constraints) + ] + + # Rename dimensions in lpcosts + lpcosts = lpcosts.rename({k: f"d({k})" for k in lpcosts.dims}) + + # Merge data + data = merge([lpcosts, *lp_constraints]) + + # An adjustment is required for lower bound constraints + for i, constraint in enumerate(constraints): + if constraint.kind == ConstraintKind.LOWER_BOUND: + data[f"b{i}"] = -data[f"b{i}"] + data[f"capacity{i}"] = -data[f"capacity{i}"] + data[f"production{i}"] = -data[f"production{i}"] + + # Ensure consistent ordering of dimensions + return data.transpose(*data.dims) + + +def selected_quantity(data: xr.Dataset, name: str) -> xr.Dataset: + """Select and rename variables for a specific quantity.""" + # Select data for the specified quantity ("capacity", "production", or "b") + result = data[[u for u in data.data_vars if str(u).startswith(name)]] + + # Rename variables ("costs" for the costs variable, 0/1/2 etc. for constraints) + return result.rename( + { + k: ("costs" if k == name else int(str(k).replace(name, ""))) + for k in result.data_vars + } + ) + + +def reshape_constraint_matrix(matrix: xr.DataArray) -> np.ndarray: + """Convert constraints matrix to a 2D np array. + + The rows of the constraints matrix will represent the constraints, and the + columns will represent the decision variables. + """ + # Before building LP we need to sort dimensions for consistency + if list(matrix.dims) != sorted(matrix.dims): + matrix = matrix.transpose(*sorted(matrix.dims)) + + # Size of the first dimension + # This dimension represents the number of constraints + size = np.prod([matrix[u].shape[0] for u in matrix.dims if str(u).startswith("c")]) + + # Reshape into a 2D array: N constraints x N decision variables + return matrix.values.reshape((size, -1)) + + +def extract_constraint_matrices( + capacities: xr.Dataset, + productions: xr.Dataset, + bs: xr.Dataset, + constraints, + *kinds: ConstraintKind, +): + """Extracts A and b matrices for constraints of specified kinds.""" + # Get indices of constraints of the specified kind + indices = [i for i in range(len(bs)) if constraints[i].kind in kinds] + + # Convert constraints matrices to 2d np arrays + capa_constraints = [reshape_constraint_matrix(capacities[i]) for i in indices] + prod_constraints = [reshape_constraint_matrix(productions[i]) for i in indices] + + # Convert constraints vectors to 1d + constraints_vectors = [bs[i].stack(constraint=sorted(bs[i].dims)) for i in indices] + + # Concatenate constraints + if capa_constraints: + A = np.concatenate( + ( + np.concatenate(capa_constraints, axis=0), + np.concatenate(prod_constraints, axis=0), + ), + axis=1, + ) + b = np.concatenate(constraints_vectors, axis=0) + else: + # If there are no constraints of the given kind, return None + A = None + b = None + return A, b + + +def back_to_muse_quantity(x: np.ndarray, template: xr.DataArray) -> xr.DataArray: + """Convert a vector of decision variables to a DataArray.""" + # First create a multidimensional dataarray based on the template + result = xr.DataArray( + x.reshape(template.shape), coords=template.coords, dims=template.dims + ) + + # Then rename the dimensions (e.g. "d(asset)" -> "asset") + return result.rename({k: str(k)[2:-1] for k in result.dims}) + + +@dataclass +class ScipyAdapter: + """Adapts MUSE data structures to scipy's linear programming solver format. + + This adapter converts data (costs and constraints) from xarray DataArrays to + the format required by scipy's linear programming solver, and back. + """ + + c: np.ndarray + capacity_template: xr.DataArray + production_template: xr.DataArray + bounds: tuple[float | None, float | None] = (0, np.inf) + A_ub: np.ndarray | None = None + b_ub: np.ndarray | None = None + A_eq: np.ndarray | None = None + b_eq: np.ndarray | None = None + + @classmethod + def from_muse_data( + cls, + costs: xr.DataArray, + constraints: list, + commodities: list[str], + timeslice_level: str | None = None, + ) -> ScipyAdapter: + """Creates a ScipyAdapter from MUSE data structures.""" + # Calculate costs for the linear problem + lpcosts = lp_costs( + capacity_costs=costs, + commodities=commodities, + timeslice_level=timeslice_level, + ) + + # Create dataset from costs and constraints + data = unified_dataset(lpcosts, *constraints) + + # Get capacity constraint matrix / costs + capacities = selected_quantity(data, "capacity") + + # Get production constraint matrix / costs + productions = selected_quantity(data, "production") + + # Get constraint vector + bs = selected_quantity(data, "b") + + # Create costs vector by concatenating capacity and production costs + c = np.concatenate( + ( + capacities["costs"].values.flatten(), + productions["costs"].values.flatten(), + ), + axis=0, + ) + + # Extract A and b for inequality constraints + A_ub, b_ub = extract_constraint_matrices( + capacities, + productions, + bs, + constraints, + ConstraintKind.UPPER_BOUND, + ConstraintKind.LOWER_BOUND, + ) + + # Extract A and b for equality constraints + A_eq, b_eq = extract_constraint_matrices( + capacities, productions, bs, constraints, ConstraintKind.EQUALITY + ) + + return cls( + c=c, + A_ub=A_ub, + b_ub=b_ub, + A_eq=A_eq, + b_eq=b_eq, + capacity_template=capacities["costs"], + production_template=productions["costs"], + ) + + def to_muse(self, x: np.ndarray) -> xr.Dataset: + """Convert scipy solver output back to MUSE format.""" + n_capa = self.capacity_template.size + capa = back_to_muse_quantity(x[:n_capa], self.capacity_template) + prod = back_to_muse_quantity(x[n_capa:], self.production_template) + return xr.Dataset({"capacity": capa, "production": prod}) + + @property + def kwargs(self): + """Dictionary of kwargs for scipy.optimize.linprog.""" + return { + "c": self.c, + "A_eq": self.A_eq, + "b_eq": self.b_eq, + "A_ub": self.A_ub, + "b_ub": self.b_ub, + "bounds": self.bounds, + } + + +def lp_costs( + capacity_costs: xr.DataArray, + commodities: list[str], + timeslice_level: str | None = None, +) -> xr.Dataset: + """Creates dataset of costs for solving with scipy's LP solver.""" + assert set(capacity_costs.dims) == {"asset", "replacement"} + + # Start with capacity costs as template (defines "asset" and "replacement" dims) + production_costs = xr.zeros_like(capacity_costs) + + # Add a "timeslice" dimension, convert multiindex to single index + production_costs = broadcast_timeslice(production_costs, level=timeslice_level) + production_costs = drop_timeslice(production_costs) + production_costs["timeslice"] = pd.Index( + production_costs.get_index("timeslice"), tupleize_cols=False + ) + + # Add a "commodity" dimension + production_costs = production_costs.expand_dims(commodity=commodities) + assert set(production_costs.dims) == { + "asset", + "replacement", + "commodity", + "timeslice", + } + + # Result is dataset of provided capacity costs and zero production costs + return xr.Dataset(dict(capacity=capacity_costs, production=production_costs)) + + +def lp_constraint(constraint, lpcosts: xr.Dataset) -> xr.Dataset: + """Transforms the constraint to LP data. + + The goal is to create from ``lpcosts.capacity``, ``constraint.capacity``, and + ``constraint.b`` a 2d-matrix ``constraint`` vs ``decision variables``. + """ + constraint = constraint.copy(deep=False) + + # Deal with timeslice multiindex + if "timeslice" in constraint.dims: + constraint = drop_timeslice(constraint) + constraint["timeslice"] = pd.Index( + constraint.get_index("timeslice"), tupleize_cols=False + ) + + # Rename dimensions in b + b = constraint.b.drop_vars(set(constraint.b.coords) - set(constraint.b.dims)) + b = b.rename({k: f"c({k})" for k in b.dims}) + + # Create capacity constraint matrix + capacity = lp_constraint_matrix(constraint.b, constraint.capacity, lpcosts.capacity) + capacity = capacity.drop_vars(set(capacity.coords) - set(capacity.dims)) + + # Create production constraint matrix + production = lp_constraint_matrix( + constraint.b, constraint.production, lpcosts.production + ) + production = production.drop_vars(set(production.coords) - set(production.dims)) + + # Combine data + result = xr.Dataset( + {"b": b, "capacity": capacity, "production": production}, attrs=constraint.attrs + ) + return result + + +def lp_constraint_matrix( + b: xr.DataArray, constraint: xr.DataArray, lpcosts: xr.DataArray +): + """Transforms one constraint block into an lp matrix.""" + from functools import reduce + + from numpy import eye + + # Sum over all dimensions that are not in the constraint or the decision variables + result = constraint.sum(set(constraint.dims) - set(lpcosts.dims) - set(b.dims)) + + # Rename dimensions for decision variables + result = result.rename( + {k: f"d({k})" for k in set(result.dims).intersection(lpcosts.dims)} + ) + + # Rename dimensions for constraints + result = result.rename( + {k: f"c({k})" for k in set(result.dims).intersection(b.dims)} + ) + + # Expand dimensions that are in the decision variables but not in the constraint + expand = set(lpcosts.dims) - set(constraint.dims) - set(b.dims) + result = result.expand_dims( + {f"d({k})": lpcosts[k].rename({k: f"d({k})"}).set_index() for k in expand} + ) + + # Expand dimensions that are in the constraint but not in the decision variables + expand = set(b.dims) - set(constraint.dims) - set(lpcosts.dims) + result = result.expand_dims( + {f"c({k})": b[k].rename({k: f"c({k})"}).set_index() for k in expand} + ) + + # Dimensions that are in both the decision variables and the constraint + diag_dims = set(b.dims).intersection(lpcosts.dims) + diag_dims = sorted(diag_dims) + + if diag_dims: + + def get_dimension(dim): + if dim in b.dims: + return b[dim].values + if dim in lpcosts.dims: + return lpcosts[dim].values + return constraint[dim].values + + diagonal_submats = [ + xr.DataArray( + eye(len(b[k])), + coords={f"c({k})": get_dimension(k), f"d({k})": get_dimension(k)}, + dims=(f"c({k})", f"d({k})"), + ) + for k in diag_dims + ] + reduced = reduce(xr.DataArray.__mul__, diagonal_submats) + if "d(timeslice)" in reduced.dims: + reduced = reduced.drop_vars("d(timeslice)") + result = result * reduced + + return result From e43f4ead2d230c4e1d16c58f64be00b620849ed9 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Thu, 5 Jun 2025 22:28:15 +0100 Subject: [PATCH 28/33] Tidy test_unmet_forecast_demand --- tests/test_demand_share.py | 61 ++++++++++---------------------------- 1 file changed, 16 insertions(+), 45 deletions(-) diff --git a/tests/test_demand_share.py b/tests/test_demand_share.py index 36b3e55d..7fbd661f 100644 --- a/tests/test_demand_share.py +++ b/tests/test_demand_share.py @@ -437,64 +437,35 @@ def test_unmet_forecast_demand(_technologies, timeslice, stock): """Test unmet forecast demand calculations. Tests three scenarios: - 1. Fully met demand - 2. Excess capacity - 3. Insufficient capacity + 1. Fully met demand - agents have exact capacity to meet demand + 2. Excess capacity - agents have more capacity than needed + 3. Insufficient capacity - agents have less capacity than needed """ - from dataclasses import dataclass - from muse.commodities import is_enduse from muse.demand_share import unmet_forecasted_demand - from muse.utilities import broadcast_over_assets - - asia_stock = stock.where(stock.region == "ASEAN", drop=True) - usa_stock = stock.where(stock.region == "USA", drop=True) - asia_market = _matching_market( - broadcast_over_assets(_technologies, asia_stock), asia_stock.capacity - ) - usa_market = _matching_market( - broadcast_over_assets(_technologies, usa_stock), usa_stock.capacity - ) - market = xr.concat((asia_market, usa_market), dim="region") - - @dataclass - class Agent: - assets: xr.Dataset + # Setup market data + market, asia_stock, usa_stock = create_regional_market(_technologies, stock) - # Test fully met demand - agents = [ - Agent(0.3 * usa_stock), - Agent(0.7 * usa_stock), - Agent(asia_stock), - ] + # Test scenario 1: Fully met demand + agents = create_test_agents(usa_stock, asia_stock) result = unmet_forecasted_demand(agents, market.consumption, _technologies) assert set(result.dims) == set(market.consumption.dims) - {"year"} assert result.values == approx(0) - # Test excess capacity - agents = [ - Agent(0.4 * usa_stock), - Agent(0.8 * usa_stock), - Agent(1.1 * asia_stock), - ] - result = unmet_forecasted_demand( - agents, - market.consumption, - _technologies, - ) + # Test scenario 2: Excess capacity (120% capacity) + agents = create_test_agents(1.2 * usa_stock, 1.2 * asia_stock) + result = unmet_forecasted_demand(agents, market.consumption, _technologies) assert set(result.dims) == set(market.consumption.dims) - {"year"} assert result.values == approx(0) - # Test insufficient capacity - agents = [ - Agent(0.5 * usa_stock), - Agent(0.5 * asia_stock), - ] + # Test scenario 3: Insufficient capacity (50% capacity) + agents = create_test_agents(0.5 * usa_stock, 0.5 * asia_stock) result = unmet_forecasted_demand(agents, market.consumption, _technologies) - comm_usage = _technologies.comm_usage.sel(commodity=market.commodity) - enduse = is_enduse(comm_usage) - assert (result.commodity == comm_usage.commodity).all() + + # Verify results for insufficient capacity + enduse = is_enduse(_technologies.comm_usage.sel(commodity=market.commodity)) + assert (result.commodity == market.commodity).all() assert result.sel(commodity=~enduse).values == approx(0) assert result.sel(commodity=enduse).values == approx( 0.5 * market.consumption.sel(commodity=enduse, year=2030).values From 7574eb318addd0691d7acbc8fb861e2b6eeda308 Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Mon, 9 Jun 2025 17:19:30 +0100 Subject: [PATCH 29/33] Combine fixtures in test_constraints --- tests/test_constraints.py | 476 +++++++++++++++++++++----------------- 1 file changed, 262 insertions(+), 214 deletions(-) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index a55c232a..75201c3a 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -3,183 +3,195 @@ import numpy as np import pandas as pd import xarray as xr -from pytest import approx, fixture, raises +from pytest import approx, fixture, mark, raises +from muse import examples +from muse.constraints import ( + demand, + demand_limiting_capacity, + lp_costs, + max_capacity_expansion, + max_production, +) +from muse.quantities import maximum_production from muse.timeslices import drop_timeslice -from muse.utilities import interpolate_capacity, reduce_assets +from muse.utilities import broadcast_over_assets, interpolate_capacity, reduce_assets CURRENT_YEAR = 2020 INVESTMENT_YEAR = 2025 -@fixture -def model(): - return "medium" - - -@fixture -def residential(model): - from muse import examples - - return examples.sector("residential", model=model) - - -@fixture -def technologies(residential): - return residential.technologies.squeeze("region").sel(year=INVESTMENT_YEAR) - - -@fixture -def search_space(model, assets): - from muse import examples - - space = examples.search_space("residential", model) - return space.sel(asset=assets.technology.values) - - -@fixture -def costs(search_space): - shape = search_space.shape - return search_space * np.arange(np.prod(shape)).reshape(shape) - - -@fixture -def assets(residential): - return next(a.assets for a in residential.agents) - - -@fixture -def capacity(assets): - return interpolate_capacity( +@fixture(params=["timeslice_as_list", "timeslice_as_multindex"]) +def constraint_data(request): + """Creates the complete dataset needed for constraint testing. + + The transformation follows these steps: + 1. Load residential sector data + 2. Extract technologies and assets + 3. Calculate capacity and market demand + 4. Create search space and costs + 5. Generate all constraint-specific data + + Returns: + dict: Contains all necessary data for constraint testing: + - technologies: Technology parameters + - assets: Asset data + - capacity: Interpolated capacity data + - market_demand: Calculated market demand + - search_space: Search space for investments + - costs: Cost data for each option + - commodities: List of relevant commodities + - lp_costs: Linear programming cost data + - constraints: Dict of all constraint matrices + """ + # Step 1: Load residential sector data + residential = examples.sector("residential", model="medium") + + # Step 2: Extract technologies and assets + technologies = residential.technologies.squeeze("region").sel(year=INVESTMENT_YEAR) + assets = next(a.assets for a in residential.agents) + + # Step 3: Calculate capacity and market demand + capacity = interpolate_capacity( reduce_assets(assets.capacity, coords=("technology", "region")), year=[CURRENT_YEAR, INVESTMENT_YEAR], ) - -@fixture -def market_demand(assets, technologies): - from muse.quantities import maximum_production - from muse.utilities import broadcast_over_assets - - # Set demand just below the maximum production of existing assets - res = 0.8 * maximum_production( + # Create initial market demand as 80% of maximum production + market_demand = 0.8 * maximum_production( broadcast_over_assets(technologies, assets), assets.capacity, ).sel(year=INVESTMENT_YEAR).groupby("technology").sum("asset").rename( technology="asset" ) - # Remove un-demanded commodities - res = res.sel(commodity=(res > 0).any(dim=["timeslice", "asset"])) - return res - - -@fixture -def commodities(market_demand): - return list(market_demand.commodity.values) - - -def test_fixtures(technologies, search_space, costs, assets, capacity, market_demand): - assert set(technologies.dims) == {"technology", "commodity"} - assert set(search_space.dims) == {"asset", "replacement"} - assert set(costs.dims) == {"asset", "replacement"} - assert set(assets.dims) == {"asset", "year"} - assert set(capacity.dims) == {"asset", "year"} - assert set(market_demand.dims) == {"asset", "commodity", "timeslice"} - - -@fixture -def lpcosts(costs, commodities): - from muse.constraints import lp_costs - - return lp_costs(costs, commodities=commodities) - - -@fixture -def max_production(market_demand, capacity, search_space, technologies): - from muse.constraints import max_production - - return max_production(market_demand, capacity, search_space, technologies) + market_demand = market_demand.sel( + commodity=(market_demand > 0).any(dim=["timeslice", "asset"]) + ) + # Step 4: Create search space and costs + search_space = examples.search_space("residential", "medium") + search_space = search_space.sel(asset=assets.technology.values) + shape = search_space.shape + costs = search_space * np.arange(np.prod(shape)).reshape(shape) + + # Get list of commodities + commodities = list(market_demand.commodity.values) + + # Step 5: Generate constraints + constraints = { + "max_production": max_production( + market_demand, capacity, search_space, technologies + ), + "demand": demand(market_demand, capacity, search_space, technologies), + "max_capacity_expansion": max_capacity_expansion( + market_demand, capacity, search_space, technologies + ), + "demand_limiting_capacity": demand_limiting_capacity( + market_demand, capacity, search_space, technologies + ), + } -@fixture -def demand_constraint(market_demand, capacity, search_space, technologies): - from muse.constraints import demand + # Testing two different timeslicing formats + if request.param == "timeslice_as_multindex": + constraints = {key: _as_list(cs) for key, cs in constraints.items()} + + # Step 6: Generate lp costs + lp_cost_data = lp_costs(costs, commodities=commodities) + + return { + "technologies": technologies, + "capacity": capacity, + "market_demand": market_demand, + "search_space": search_space, + "costs": costs, + "commodities": commodities, + "lp_costs": lp_cost_data, + "constraints": constraints, + } - return demand(market_demand, capacity, search_space, technologies) +def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Dataset]: + """Helper function to convert timeslice data to list format.""" + if "timeslice" in data.dims: + data = data.copy(deep=False) + index = pd.MultiIndex.from_tuples( + data.get_index("timeslice"), names=("month", "day", "hour") + ) + mindex_coords = xr.Coordinates.from_pandas_multiindex(index, "timeslice") + data = drop_timeslice(data).assign_coords(mindex_coords) + return data -@fixture -def max_capacity_expansion(market_demand, capacity, search_space, technologies): - from muse.constraints import max_capacity_expansion - return max_capacity_expansion(market_demand, capacity, search_space, technologies) +def test_fixtures(constraint_data): + """Validating that the fixture data has appropriate dimensions.""" + assert set(constraint_data["technologies"].dims) == {"technology", "commodity"} + assert set(constraint_data["search_space"].dims) == {"asset", "replacement"} + assert set(constraint_data["costs"].dims) == {"asset", "replacement"} + assert set(constraint_data["capacity"].dims) == {"asset", "year"} + assert set(constraint_data["market_demand"].dims) == { + "asset", + "commodity", + "timeslice", + } -@fixture -def demand_limiting_capacity(market_demand, capacity, search_space, technologies): - from muse.constraints import demand_limiting_capacity +@mark.usefixtures("save_registries") +def test_objective_registration(): + from muse.objectives import OBJECTIVES, register_objective - return demand_limiting_capacity(market_demand, capacity, search_space, technologies) + @register_objective + def a_objective(*args, **kwargs): + pass + assert "a_objective" in OBJECTIVES + assert OBJECTIVES["a_objective"] is a_objective -@fixture -def constraint(max_production): - return max_production + @register_objective(name="something") + def b_objective(*args, **kwargs): + pass + assert "something" in OBJECTIVES + assert OBJECTIVES["something"] is b_objective -@fixture(params=["timeslice_as_list", "timeslice_as_multindex"]) -def constraints( - request, - max_production, - demand_constraint, - demand_limiting_capacity, - max_capacity_expansion, -): - constraints = [ - max_production, - demand_limiting_capacity, - demand_constraint, - max_capacity_expansion, - ] - if request.param == "timeslice_as_multindex": - constraints = [_as_list(cs) for cs in constraints] - return constraints +def test_constraints_dimensions(constraint_data): + """Test dimensions of all constraint matrices.""" + constraints = constraint_data["constraints"] -def test_constraints_dimensions( - max_production, demand_constraint, demand_limiting_capacity, max_capacity_expansion -): # Max production constraint max_prod_dims = {"asset", "commodity", "replacement", "timeslice"} - assert set(max_production.capacity.dims) == max_prod_dims - assert set(max_production.production.dims) == max_prod_dims - assert set(max_production.b.dims) == max_prod_dims + assert set(constraints["max_production"].capacity.dims) == max_prod_dims + assert set(constraints["max_production"].production.dims) == max_prod_dims + assert set(constraints["max_production"].b.dims) == max_prod_dims # Demand constraint - assert set(demand_constraint.capacity.dims) == set() - assert set(demand_constraint.production.dims) == set() - assert set(demand_constraint.b.dims) == {"asset", "commodity", "timeslice"} + assert set(constraints["demand"].capacity.dims) == set() + assert set(constraints["demand"].production.dims) == set() + assert set(constraints["demand"].b.dims) == {"asset", "commodity", "timeslice"} # Demand limiting capacity constraint - assert set(demand_limiting_capacity.capacity.dims) == { + assert set(constraints["demand_limiting_capacity"].capacity.dims) == { "asset", "commodity", "replacement", } - assert set(demand_limiting_capacity.production.dims) == set() - assert set(demand_limiting_capacity.b.dims) == {"asset", "commodity"} + assert set(constraints["demand_limiting_capacity"].production.dims) == set() + assert set(constraints["demand_limiting_capacity"].b.dims) == {"asset", "commodity"} # Max capacity expansion constraint - assert set(max_capacity_expansion.capacity.dims) == set() - assert set(max_capacity_expansion.production.dims) == set() - assert set(max_capacity_expansion.b.dims) == {"replacement"} + assert set(constraints["max_capacity_expansion"].capacity.dims) == set() + assert set(constraints["max_capacity_expansion"].production.dims) == set() + assert set(constraints["max_capacity_expansion"].b.dims) == {"replacement"} -def test_lp_constraints_matrix_b_is_scalar(constraint, lpcosts): +def test_lp_constraints_matrix_b_is_scalar(constraint_data): """B is a scalar - output should be equivalent to a single row matrix.""" from muse.constraints import lp_constraint_matrix + constraint = constraint_data["constraints"]["max_production"] + lpcosts = constraint_data["lp_costs"] + for attr in ["capacity", "production"]: lpconstraint = lp_constraint_matrix( xr.DataArray(1), getattr(constraint, attr), getattr(lpcosts, attr) @@ -191,10 +203,13 @@ def test_lp_constraints_matrix_b_is_scalar(constraint, lpcosts): } -def test_max_production_constraint_diagonal(constraint, lpcosts): +def test_max_production_constraint_diagonal(constraint_data): """Test production side of max capacity production is diagonal.""" from muse.constraints import lp_constraint_matrix + constraint = constraint_data["constraints"]["max_production"] + lpcosts = constraint_data["lp_costs"] + # Test capacity constraints result = lp_constraint_matrix(constraint.b, constraint.capacity, lpcosts.capacity) decision_dims = {f"d({x})" for x in lpcosts.capacity.dims} @@ -219,9 +234,12 @@ def test_max_production_constraint_diagonal(constraint, lpcosts): assert stacked.values == approx(np.eye(stacked.shape[0])) -def test_lp_constraint(constraint, lpcosts): +def test_lp_constraint(constraint_data): from muse.constraints import lp_constraint + constraint = constraint_data["constraints"]["max_production"] + lpcosts = constraint_data["lp_costs"] + result = lp_constraint(constraint, lpcosts) constraint_dims = { f"c({x})" for x in set(lpcosts.production.dims).union(constraint.b.dims) @@ -244,13 +262,17 @@ def test_lp_constraint(constraint, lpcosts): assert result.b.values == approx(0) -def test_to_scipy_adapter_maxprod(costs, max_production, commodities, lpcosts): +def test_to_scipy_adapter_maxprod(constraint_data): """Test scipy adapter with max production constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - costs, constraints=[max_production], commodities=commodities + constraint_data["costs"], + constraints=[constraint_data["constraints"]["max_production"]], + commodities=constraint_data["commodities"], ) + lpcosts = constraint_data["lp_costs"] + assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) assert adapter.A_ub is not None @@ -280,13 +302,17 @@ def test_to_scipy_adapter_maxprod(costs, max_production, commodities, lpcosts): assert adapter.A_ub[:, capsize:] == approx(np.eye(prodsize)) -def test_to_scipy_adapter_demand(costs, demand_constraint, commodities, lpcosts): +def test_to_scipy_adapter_demand(constraint_data): """Test scipy adapter with demand constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - costs, constraints=[demand_constraint], commodities=commodities + constraint_data["costs"], + constraints=[constraint_data["constraints"]["demand"]], + commodities=constraint_data["commodities"], ) + lpcosts = constraint_data["lp_costs"] + assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) assert adapter.A_ub is not None @@ -310,15 +336,17 @@ def test_to_scipy_adapter_demand(costs, demand_constraint, commodities, lpcosts) assert set(adapter.A_ub[:, capsize:].flatten()) == {0.0, -1.0} -def test_to_scipy_adapter_max_capacity_expansion( - costs, max_capacity_expansion, commodities, lpcosts -): +def test_to_scipy_adapter_max_capacity_expansion(constraint_data): """Test scipy adapter with max capacity expansion constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - costs, constraints=[max_capacity_expansion], commodities=commodities + constraint_data["costs"], + constraints=[constraint_data["constraints"]["max_capacity_expansion"]], + commodities=constraint_data["commodities"], ) + lpcosts = constraint_data["lp_costs"] + assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) assert adapter.A_ub is not None @@ -342,10 +370,16 @@ def test_to_scipy_adapter_max_capacity_expansion( assert set(adapter.A_ub[:, :capsize].flatten()) == {0.0, 1.0} -def test_scipy_adapter_no_constraint(costs, commodities, lpcosts): +def test_scipy_adapter_no_constraint(constraint_data): from muse.constraints import ScipyAdapter - adapter = ScipyAdapter.factory(costs, constraints=[], commodities=commodities) + adapter = ScipyAdapter.factory( + constraint_data["costs"], + constraints=[], + commodities=constraint_data["commodities"], + ) + lpcosts = constraint_data["lp_costs"] + assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) assert all( @@ -355,9 +389,10 @@ def test_scipy_adapter_no_constraint(costs, commodities, lpcosts): assert adapter.c.size == lpcosts.capacity.size + lpcosts.production.size -def test_back_to_muse_quantities(lpcosts): +def test_back_to_muse_quantities(constraint_data): from muse.constraints import ScipyAdapter + lpcosts = constraint_data["lp_costs"] data = ScipyAdapter._unified_dataset(lpcosts) # Test capacity @@ -382,9 +417,10 @@ def test_back_to_muse_quantities(lpcosts): assert (copy == lpcosts.production).all() -def test_back_to_muse_all(lpcosts): +def test_back_to_muse_all(constraint_data): from muse.constraints import ScipyAdapter + lpcosts = constraint_data["lp_costs"] data = ScipyAdapter._unified_dataset(lpcosts) lpcapacity = ScipyAdapter._selected_quantity(data, "capacity") lpproduction = ScipyAdapter._selected_quantity(data, "production") @@ -408,94 +444,87 @@ def test_back_to_muse_all(lpcosts): assert (copy.production == lpcosts.production).all() -def test_scipy_adapter_back_to_muse(costs, constraints, commodities, lpcosts): - """Test converting back from scipy adapter format to MUSE format.""" +def test_scipy_adapter_standard_constraints(constraint_data): from muse.constraints import ScipyAdapter - data = ScipyAdapter._unified_dataset(lpcosts) - lpcapacity = ScipyAdapter._selected_quantity(data, "capacity") - lpproduction = ScipyAdapter._selected_quantity(data, "production") + constraints = [ + constraint_data["constraints"]["max_production"], + constraint_data["constraints"]["max_capacity_expansion"], + constraint_data["constraints"]["demand"], + constraint_data["constraints"]["demand_limiting_capacity"], + ] - x = np.concatenate( - ( - lpcosts.capacity.transpose( - *[u[2:-1] for u in lpcapacity.dims] - ).values.flatten(), - lpcosts.production.transpose( - *[u[2:-1] for u in lpproduction.dims] - ).values.flatten(), - ) + adapter = ScipyAdapter.factory( + constraint_data["costs"], + constraints, + commodities=constraint_data["commodities"], ) - adapter = ScipyAdapter.factory(costs, constraints, commodities=commodities) - assert (adapter.to_muse(x).capacity == lpcosts.capacity).all() - assert (adapter.to_muse(x).production == lpcosts.production).all() - - -def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Dataset]: - if "timeslice" in data.dims: - data = data.copy(deep=False) - index = pd.MultiIndex.from_tuples( - data.get_index("timeslice"), names=("month", "day", "hour") - ) - mindex_coords = xr.Coordinates.from_pandas_multiindex(index, "timeslice") - data = drop_timeslice(data).assign_coords(mindex_coords) - return data - - -def test_scipy_adapter_standard_constraints(costs, constraints, commodities): - from muse.constraints import ScipyAdapter - - adapter = ScipyAdapter.factory(costs, constraints, commodities=commodities) - constraint_map = {cs.name: cs for cs in constraints} - maxprod = constraint_map["max_production"] - maxcapa = constraint_map["max capacity expansion"] - demand = constraint_map["demand"] - dlc = constraint_map["demand_limiting_capacity"] - n_constraints = adapter.b_ub.size n_decision_vars = adapter.c.size - assert n_decision_vars == costs.size + maxprod.production.size + assert ( + n_decision_vars + == constraint_data["costs"].size + constraints[0].production.size + ) assert adapter.b_eq is None assert adapter.A_eq is None assert adapter.A_ub.shape == (n_constraints, n_decision_vars) - assert n_constraints == sum(c.b.size for c in [demand, maxprod, maxcapa, dlc]) + assert n_constraints == sum(c.b.size for c in constraints) -def test_scipy_solver(technologies, costs, constraints, commodities): +def test_scipy_solver(constraint_data): """Test the scipy solver for demand matching.""" from muse.investments import scipy_match_demand + constraints = [ + constraint_data["constraints"]["max_production"], + constraint_data["constraints"]["max_capacity_expansion"], + constraint_data["constraints"]["demand"], + constraint_data["constraints"]["demand_limiting_capacity"], + ] + solution = scipy_match_demand( - costs=costs, - search_space=xr.ones_like(costs), - technologies=technologies, + costs=constraint_data["costs"], + search_space=xr.ones_like(constraint_data["costs"]), + technologies=constraint_data["technologies"], constraints=constraints, - commodities=commodities, + commodities=constraint_data["commodities"], ) assert isinstance(solution, xr.DataArray) assert set(solution.dims) == {"asset", "replacement"} -def test_minimum_service( - market_demand, capacity, search_space, technologies, constraints -): +def test_minimum_service(constraint_data): from muse.constraints import minimum_service # Test with no minimum service factor - assert minimum_service(market_demand, capacity, search_space, technologies) is None + assert ( + minimum_service( + constraint_data["market_demand"], + constraint_data["capacity"], + constraint_data["search_space"], + constraint_data["technologies"], + ) + is None + ) # Test with minimum service factor + technologies = constraint_data["technologies"].copy() technologies["minimum_service_factor"] = 0.4 * xr.ones_like( technologies.technology, dtype=float ) - min_service = minimum_service(market_demand, capacity, search_space, technologies) - constraints.append(min_service) + min_service = minimum_service( + constraint_data["market_demand"], + constraint_data["capacity"], + constraint_data["search_space"], + technologies, + ) assert isinstance(min_service, xr.Dataset) -def test_max_capacity_expansion_properties(max_capacity_expansion): +def test_max_capacity_expansion_properties(constraint_data): + max_capacity_expansion = constraint_data["constraints"]["max_capacity_expansion"] assert (max_capacity_expansion.capacity == 1).all() assert max_capacity_expansion.production == 0 assert max_capacity_expansion.b.dims == ("replacement",) @@ -506,32 +535,43 @@ def test_max_capacity_expansion_properties(max_capacity_expansion): ).all() -def test_max_capacity_expansion_no_limits( - market_demand, capacity, search_space, technologies -): +def test_max_capacity_expansion_no_limits(constraint_data): from muse.constraints import max_capacity_expansion - techs = technologies.drop_vars( + techs = constraint_data["technologies"].drop_vars( ["max_capacity_addition", "max_capacity_growth", "total_capacity_limit"] ) - assert max_capacity_expansion(market_demand, capacity, search_space, techs) is None + assert ( + max_capacity_expansion( + constraint_data["market_demand"], + constraint_data["capacity"], + constraint_data["search_space"], + techs, + ) + is None + ) -def test_max_capacity_expansion_seed( - market_demand, capacity, search_space, technologies -): +def test_max_capacity_expansion_seed(constraint_data): from muse.constraints import max_capacity_expansion seed = 10 + technologies = constraint_data["technologies"].copy() technologies["growth_seed"] = seed # Test different capacity scenarios scenarios = [0, seed, 2 * seed] results = [] for cap in scenarios: + capacity = constraint_data["capacity"].copy() capacity.sel(year=2020)[:] = cap results.append( - max_capacity_expansion(market_demand, capacity, search_space, technologies) + max_capacity_expansion( + constraint_data["market_demand"], + capacity, + constraint_data["search_space"], + technologies, + ) ) # Zero capacity should match seed capacity @@ -540,11 +580,10 @@ def test_max_capacity_expansion_seed( assert results[0].b.values != approx(results[2].b.values) -def test_max_capacity_expansion_infinite_limits( - market_demand, capacity, search_space, technologies -): +def test_max_capacity_expansion_infinite_limits(constraint_data): from muse.constraints import max_capacity_expansion + technologies = constraint_data["technologies"].copy() for limit in [ "max_capacity_addition", "max_capacity_growth", @@ -552,16 +591,25 @@ def test_max_capacity_expansion_infinite_limits( ]: technologies[limit] = np.inf with raises(ValueError): - max_capacity_expansion(market_demand, capacity, search_space, technologies) + max_capacity_expansion( + constraint_data["market_demand"], + constraint_data["capacity"], + constraint_data["search_space"], + technologies, + ) + +def test_max_production(constraint_data): + assert (constraint_data["constraints"]["max_production"].capacity <= 0).all() -def test_max_production(max_production): - assert (max_production.capacity <= 0).all() +def test_demand_limiting_capacity(constraint_data): + demand_limiting_capacity = constraint_data["constraints"][ + "demand_limiting_capacity" + ] + max_production = constraint_data["constraints"]["max_production"] + demand_constraint = constraint_data["constraints"]["demand"] -def test_demand_limiting_capacity( - demand_limiting_capacity, max_production, demand_constraint -): # Test capacity values expected_capacity = ( -max_production.capacity.max("timeslice").values From dcc4f51b322e0cd562405f94feb5ca71be99e62c Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Mon, 9 Jun 2025 20:10:54 +0100 Subject: [PATCH 30/33] Further tidy test_constraints fixtures --- tests/test_constraints.py | 288 +++++++++++++++----------------------- 1 file changed, 112 insertions(+), 176 deletions(-) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index 75201c3a..b3902c06 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -21,37 +21,25 @@ INVESTMENT_YEAR = 2025 -@fixture(params=["timeslice_as_list", "timeslice_as_multindex"]) -def constraint_data(request): - """Creates the complete dataset needed for constraint testing. - - The transformation follows these steps: - 1. Load residential sector data - 2. Extract technologies and assets - 3. Calculate capacity and market demand - 4. Create search space and costs - 5. Generate all constraint-specific data +@fixture +def model_data(): + """Model data required for the constraints. Returns: - dict: Contains all necessary data for constraint testing: - - technologies: Technology parameters - - assets: Asset data - - capacity: Interpolated capacity data - - market_demand: Calculated market demand - - search_space: Search space for investments - - costs: Cost data for each option - - commodities: List of relevant commodities - - lp_costs: Linear programming cost data - - constraints: Dict of all constraint matrices + dict: Contains all necessary data for building constraints: + - technologies: Technology data for a single year + - capacity: Capacity data for assets in the current and investment year + - demand: Demand for the investment year + - search_space: Search space for the assets """ - # Step 1: Load residential sector data + # Load residential sector data residential = examples.sector("residential", model="medium") - # Step 2: Extract technologies and assets + # Extract technologies and assets technologies = residential.technologies.squeeze("region").sel(year=INVESTMENT_YEAR) assets = next(a.assets for a in residential.agents) - # Step 3: Calculate capacity and market demand + # Calculate capacity and market demand capacity = interpolate_capacity( reduce_assets(assets.capacity, coords=("technology", "region")), year=[CURRENT_YEAR, INVESTMENT_YEAR], @@ -69,48 +57,32 @@ def constraint_data(request): commodity=(market_demand > 0).any(dim=["timeslice", "asset"]) ) - # Step 4: Create search space and costs + # Create search space search_space = examples.search_space("residential", "medium") search_space = search_space.sel(asset=assets.technology.values) - shape = search_space.shape - costs = search_space * np.arange(np.prod(shape)).reshape(shape) - - # Get list of commodities - commodities = list(market_demand.commodity.values) - - # Step 5: Generate constraints - constraints = { - "max_production": max_production( - market_demand, capacity, search_space, technologies - ), - "demand": demand(market_demand, capacity, search_space, technologies), - "max_capacity_expansion": max_capacity_expansion( - market_demand, capacity, search_space, technologies - ), - "demand_limiting_capacity": demand_limiting_capacity( - market_demand, capacity, search_space, technologies - ), - } - - # Testing two different timeslicing formats - if request.param == "timeslice_as_multindex": - constraints = {key: _as_list(cs) for key, cs in constraints.items()} - - # Step 6: Generate lp costs - lp_cost_data = lp_costs(costs, commodities=commodities) + # Return dictionary of data return { "technologies": technologies, "capacity": capacity, - "market_demand": market_demand, + "demand": market_demand, "search_space": search_space, - "costs": costs, - "commodities": commodities, - "lp_costs": lp_cost_data, - "constraints": constraints, } +@fixture(params=["timeslice_as_list", "timeslice_as_multindex"]) +def constraints(request, model_data): + constraints = { + "max_production": max_production(**model_data), + "demand": demand(**model_data), + "max_capacity_expansion": max_capacity_expansion(**model_data), + "demand_limiting_capacity": demand_limiting_capacity(**model_data), + } + if request.param == "timeslice_as_multindex": + constraints = {key: _as_list(cs) for key, cs in constraints.items()} + return constraints + + def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Dataset]: """Helper function to convert timeslice data to list format.""" if "timeslice" in data.dims: @@ -123,13 +95,34 @@ def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Da return data -def test_fixtures(constraint_data): +@fixture +def lp_inputs(model_data): + """Inputs to the lp adapter, in addition to the constraints.""" + # Make up capacity costs data + shape = model_data["search_space"].shape + costs = model_data["search_space"] * np.arange(np.prod(shape)).reshape(shape) + + # List of commodities + commodities = list(model_data["demand"].commodity.values) + + return { + "capacity_costs": costs, + "commodities": commodities, + } + + +@fixture +def lpcosts(lp_inputs): + """Benchmark lpcosts dataset to test against.""" + return lp_costs(**lp_inputs) + + +def test_model_data(model_data): """Validating that the fixture data has appropriate dimensions.""" - assert set(constraint_data["technologies"].dims) == {"technology", "commodity"} - assert set(constraint_data["search_space"].dims) == {"asset", "replacement"} - assert set(constraint_data["costs"].dims) == {"asset", "replacement"} - assert set(constraint_data["capacity"].dims) == {"asset", "year"} - assert set(constraint_data["market_demand"].dims) == { + assert set(model_data["technologies"].dims) == {"technology", "commodity"} + assert set(model_data["search_space"].dims) == {"asset", "replacement"} + assert set(model_data["capacity"].dims) == {"asset", "year"} + assert set(model_data["demand"].dims) == { "asset", "commodity", "timeslice", @@ -155,10 +148,8 @@ def b_objective(*args, **kwargs): assert OBJECTIVES["something"] is b_objective -def test_constraints_dimensions(constraint_data): +def test_constraints_dimensions(constraints): """Test dimensions of all constraint matrices.""" - constraints = constraint_data["constraints"] - # Max production constraint max_prod_dims = {"asset", "commodity", "replacement", "timeslice"} assert set(constraints["max_production"].capacity.dims) == max_prod_dims @@ -185,12 +176,11 @@ def test_constraints_dimensions(constraint_data): assert set(constraints["max_capacity_expansion"].b.dims) == {"replacement"} -def test_lp_constraints_matrix_b_is_scalar(constraint_data): +def test_lp_constraints_matrix_b_is_scalar(lpcosts, constraints): """B is a scalar - output should be equivalent to a single row matrix.""" from muse.constraints import lp_constraint_matrix - constraint = constraint_data["constraints"]["max_production"] - lpcosts = constraint_data["lp_costs"] + constraint = constraints["max_production"] for attr in ["capacity", "production"]: lpconstraint = lp_constraint_matrix( @@ -203,12 +193,11 @@ def test_lp_constraints_matrix_b_is_scalar(constraint_data): } -def test_max_production_constraint_diagonal(constraint_data): +def test_max_production_constraint_diagonal(lpcosts, constraints): """Test production side of max capacity production is diagonal.""" from muse.constraints import lp_constraint_matrix - constraint = constraint_data["constraints"]["max_production"] - lpcosts = constraint_data["lp_costs"] + constraint = constraints["max_production"] # Test capacity constraints result = lp_constraint_matrix(constraint.b, constraint.capacity, lpcosts.capacity) @@ -234,11 +223,10 @@ def test_max_production_constraint_diagonal(constraint_data): assert stacked.values == approx(np.eye(stacked.shape[0])) -def test_lp_constraint(constraint_data): +def test_lp_constraint(lpcosts, constraints): from muse.constraints import lp_constraint - constraint = constraint_data["constraints"]["max_production"] - lpcosts = constraint_data["lp_costs"] + constraint = constraints["max_production"] result = lp_constraint(constraint, lpcosts) constraint_dims = { @@ -262,16 +250,14 @@ def test_lp_constraint(constraint_data): assert result.b.values == approx(0) -def test_to_scipy_adapter_maxprod(constraint_data): +def test_to_scipy_adapter_maxprod(lp_inputs, lpcosts, constraints): """Test scipy adapter with max production constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - constraint_data["costs"], - constraints=[constraint_data["constraints"]["max_production"]], - commodities=constraint_data["commodities"], + **lp_inputs, + constraints=[constraints["max_production"]], ) - lpcosts = constraint_data["lp_costs"] assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) @@ -302,16 +288,14 @@ def test_to_scipy_adapter_maxprod(constraint_data): assert adapter.A_ub[:, capsize:] == approx(np.eye(prodsize)) -def test_to_scipy_adapter_demand(constraint_data): +def test_to_scipy_adapter_demand(lp_inputs, lpcosts, constraints): """Test scipy adapter with demand constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - constraint_data["costs"], - constraints=[constraint_data["constraints"]["demand"]], - commodities=constraint_data["commodities"], + **lp_inputs, + constraints=[constraints["demand"]], ) - lpcosts = constraint_data["lp_costs"] assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) @@ -336,16 +320,14 @@ def test_to_scipy_adapter_demand(constraint_data): assert set(adapter.A_ub[:, capsize:].flatten()) == {0.0, -1.0} -def test_to_scipy_adapter_max_capacity_expansion(constraint_data): +def test_to_scipy_adapter_max_capacity_expansion(lp_inputs, lpcosts, constraints): """Test scipy adapter with max capacity expansion constraint.""" from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - constraint_data["costs"], - constraints=[constraint_data["constraints"]["max_capacity_expansion"]], - commodities=constraint_data["commodities"], + **lp_inputs, + constraints=[constraints["max_capacity_expansion"]], ) - lpcosts = constraint_data["lp_costs"] assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) @@ -370,15 +352,13 @@ def test_to_scipy_adapter_max_capacity_expansion(constraint_data): assert set(adapter.A_ub[:, :capsize].flatten()) == {0.0, 1.0} -def test_scipy_adapter_no_constraint(constraint_data): +def test_scipy_adapter_no_constraint(lp_inputs, lpcosts): from muse.constraints import ScipyAdapter adapter = ScipyAdapter.factory( - constraint_data["costs"], + **lp_inputs, constraints=[], - commodities=constraint_data["commodities"], ) - lpcosts = constraint_data["lp_costs"] assert set(adapter.kwargs) == {"c", "A_ub", "b_ub", "A_eq", "b_eq", "bounds"} assert adapter.bounds == (0, np.inf) @@ -389,10 +369,9 @@ def test_scipy_adapter_no_constraint(constraint_data): assert adapter.c.size == lpcosts.capacity.size + lpcosts.production.size -def test_back_to_muse_quantities(constraint_data): +def test_back_to_muse_quantities(lpcosts): from muse.constraints import ScipyAdapter - lpcosts = constraint_data["lp_costs"] data = ScipyAdapter._unified_dataset(lpcosts) # Test capacity @@ -417,10 +396,9 @@ def test_back_to_muse_quantities(constraint_data): assert (copy == lpcosts.production).all() -def test_back_to_muse_all(constraint_data): +def test_back_to_muse_all(lpcosts): from muse.constraints import ScipyAdapter - lpcosts = constraint_data["lp_costs"] data = ScipyAdapter._unified_dataset(lpcosts) lpcapacity = ScipyAdapter._selected_quantity(data, "capacity") lpproduction = ScipyAdapter._selected_quantity(data, "production") @@ -444,87 +422,57 @@ def test_back_to_muse_all(constraint_data): assert (copy.production == lpcosts.production).all() -def test_scipy_adapter_standard_constraints(constraint_data): +def test_scipy_adapter_standard_constraints(lp_inputs, constraints): from muse.constraints import ScipyAdapter - constraints = [ - constraint_data["constraints"]["max_production"], - constraint_data["constraints"]["max_capacity_expansion"], - constraint_data["constraints"]["demand"], - constraint_data["constraints"]["demand_limiting_capacity"], - ] - - adapter = ScipyAdapter.factory( - constraint_data["costs"], - constraints, - commodities=constraint_data["commodities"], - ) + adapter = ScipyAdapter.factory(**lp_inputs, constraints=constraints.values()) n_constraints = adapter.b_ub.size n_decision_vars = adapter.c.size + maxprod_constraint = constraints["max_production"] assert ( n_decision_vars - == constraint_data["costs"].size + constraints[0].production.size + == lp_inputs["capacity_costs"].size + maxprod_constraint.production.size ) assert adapter.b_eq is None assert adapter.A_eq is None assert adapter.A_ub.shape == (n_constraints, n_decision_vars) - assert n_constraints == sum(c.b.size for c in constraints) + assert n_constraints == sum(c.b.size for c in constraints.values()) -def test_scipy_solver(constraint_data): +def test_scipy_solver(model_data, lp_inputs, constraints): """Test the scipy solver for demand matching.""" from muse.investments import scipy_match_demand - constraints = [ - constraint_data["constraints"]["max_production"], - constraint_data["constraints"]["max_capacity_expansion"], - constraint_data["constraints"]["demand"], - constraint_data["constraints"]["demand_limiting_capacity"], - ] - solution = scipy_match_demand( - costs=constraint_data["costs"], - search_space=xr.ones_like(constraint_data["costs"]), - technologies=constraint_data["technologies"], - constraints=constraints, - commodities=constraint_data["commodities"], + capacity_costs=lp_inputs["capacity_costs"], + commodities=lp_inputs["commodities"], + search_space=model_data["search_space"], + technologies=model_data["technologies"], + constraints=constraints.values(), ) assert isinstance(solution, xr.DataArray) assert set(solution.dims) == {"asset", "replacement"} -def test_minimum_service(constraint_data): +def test_minimum_service(model_data): from muse.constraints import minimum_service - # Test with no minimum service factor - assert ( - minimum_service( - constraint_data["market_demand"], - constraint_data["capacity"], - constraint_data["search_space"], - constraint_data["technologies"], - ) - is None - ) + # Test with no minimum service factor (default) + assert minimum_service(**model_data) is None # Test with minimum service factor - technologies = constraint_data["technologies"].copy() + technologies = model_data["technologies"].copy() technologies["minimum_service_factor"] = 0.4 * xr.ones_like( technologies.technology, dtype=float ) - min_service = minimum_service( - constraint_data["market_demand"], - constraint_data["capacity"], - constraint_data["search_space"], - technologies, - ) + min_service = minimum_service(**{**model_data, "technologies": technologies}) assert isinstance(min_service, xr.Dataset) -def test_max_capacity_expansion_properties(constraint_data): - max_capacity_expansion = constraint_data["constraints"]["max_capacity_expansion"] +def test_max_capacity_expansion_properties(constraints): + max_capacity_expansion = constraints["max_capacity_expansion"] assert (max_capacity_expansion.capacity == 1).all() assert max_capacity_expansion.production == 0 assert max_capacity_expansion.b.dims == ("replacement",) @@ -535,42 +483,37 @@ def test_max_capacity_expansion_properties(constraint_data): ).all() -def test_max_capacity_expansion_no_limits(constraint_data): +def test_max_capacity_expansion_no_limits(model_data): from muse.constraints import max_capacity_expansion - techs = constraint_data["technologies"].drop_vars( + technologies = model_data["technologies"].drop_vars( ["max_capacity_addition", "max_capacity_growth", "total_capacity_limit"] ) assert ( - max_capacity_expansion( - constraint_data["market_demand"], - constraint_data["capacity"], - constraint_data["search_space"], - techs, - ) - is None + max_capacity_expansion(**{**model_data, "technologies": technologies}) is None ) -def test_max_capacity_expansion_seed(constraint_data): +def test_max_capacity_expansion_seed(model_data): from muse.constraints import max_capacity_expansion seed = 10 - technologies = constraint_data["technologies"].copy() + technologies = model_data["technologies"].copy() technologies["growth_seed"] = seed # Test different capacity scenarios scenarios = [0, seed, 2 * seed] results = [] for cap in scenarios: - capacity = constraint_data["capacity"].copy() + capacity = model_data["capacity"].copy() capacity.sel(year=2020)[:] = cap results.append( max_capacity_expansion( - constraint_data["market_demand"], - capacity, - constraint_data["search_space"], - technologies, + **{ + **model_data, + "technologies": technologies, + "capacity": capacity, + } ) ) @@ -580,10 +523,10 @@ def test_max_capacity_expansion_seed(constraint_data): assert results[0].b.values != approx(results[2].b.values) -def test_max_capacity_expansion_infinite_limits(constraint_data): +def test_max_capacity_expansion_infinite_limits(model_data): from muse.constraints import max_capacity_expansion - technologies = constraint_data["technologies"].copy() + technologies = model_data["technologies"].copy() for limit in [ "max_capacity_addition", "max_capacity_growth", @@ -591,24 +534,17 @@ def test_max_capacity_expansion_infinite_limits(constraint_data): ]: technologies[limit] = np.inf with raises(ValueError): - max_capacity_expansion( - constraint_data["market_demand"], - constraint_data["capacity"], - constraint_data["search_space"], - technologies, - ) + max_capacity_expansion(**{**model_data, "technologies": technologies}) -def test_max_production(constraint_data): - assert (constraint_data["constraints"]["max_production"].capacity <= 0).all() +def test_max_production(constraints): + assert (constraints["max_production"].capacity <= 0).all() -def test_demand_limiting_capacity(constraint_data): - demand_limiting_capacity = constraint_data["constraints"][ - "demand_limiting_capacity" - ] - max_production = constraint_data["constraints"]["max_production"] - demand_constraint = constraint_data["constraints"]["demand"] +def test_demand_limiting_capacity(constraints): + demand_limiting_capacity = constraints["demand_limiting_capacity"] + max_production = constraints["max_production"] + demand_constraint = constraints["demand"] # Test capacity values expected_capacity = ( From 2158b78aae770ee83cab2b059649d798519621bb Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Tue, 10 Jun 2025 10:52:50 +0100 Subject: [PATCH 31/33] Add back missing docstring details --- src/muse/lp_adapter.py | 74 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 70 insertions(+), 4 deletions(-) diff --git a/src/muse/lp_adapter.py b/src/muse/lp_adapter.py index 27e5615a..31f55e3c 100644 --- a/src/muse/lp_adapter.py +++ b/src/muse/lp_adapter.py @@ -87,7 +87,11 @@ def extract_constraint_matrices( constraints, *kinds: ConstraintKind, ): - """Extracts A and b matrices for constraints of specified kinds.""" + """Extracts A and b matrices for constraints of specified kinds. + + These will end up as A_ub and b_ub for inequality constraints, and A_eq and b_eq for + equality constraints (see ScipyAdapter). + """ # Get indices of constraints of the specified kind indices = [i for i in range(len(bs)) if constraints[i].kind in kinds] @@ -116,7 +120,13 @@ def extract_constraint_matrices( def back_to_muse_quantity(x: np.ndarray, template: xr.DataArray) -> xr.DataArray: - """Convert a vector of decision variables to a DataArray.""" + """Convert a vector of decision variables to a DataArray. + + Args: + x: 1D vector of decision variables, outputted from the scipy solver. + template: Template for the decision variables. This may be for either + capacity or production variables. + """ # First create a multidimensional dataarray based on the template result = xr.DataArray( x.reshape(template.shape), coords=template.coords, dims=template.dims @@ -230,7 +240,27 @@ def lp_costs( commodities: list[str], timeslice_level: str | None = None, ) -> xr.Dataset: - """Creates dataset of costs for solving with scipy's LP solver.""" + """Creates dataset of costs for solving with scipy's LP solver. + + Importantly, this also defines the decision variables in the linear program. + + The costs applied to the capacity decision variables are provided. This should + have dimensions "asset" and "replacement". In other words, capacity addition + is solved for each replacement technology for each existing asset. + + No cost is applied to the production decision variables. Thus, the production + component of the resulting dataset is zero, with dimensions determining the + production decision variables. This will have dimensions "asset", "replacement", + "commodity", and "timeslice". In other words, production is solved for each + replacement technology for each existing asset, for each commodity, and for each + timeslice. + + Args: + capacity_costs: DataArray with dimensions "asset" and "replacement" defining the + costs of adding capacity to the system. + commodities: List of commodities to create production decision variables for. + timeslice_level: The timeslice level of the linear problem. + """ assert set(capacity_costs.dims) == {"asset", "replacement"} # Start with capacity costs as template (defines "asset" and "replacement" dims) @@ -261,6 +291,22 @@ def lp_constraint(constraint, lpcosts: xr.Dataset) -> xr.Dataset: The goal is to create from ``lpcosts.capacity``, ``constraint.capacity``, and ``constraint.b`` a 2d-matrix ``constraint`` vs ``decision variables``. + + #. The dimensions of ``constraint.b`` are the constraint dimensions. They are + renamed ``"c(xxx)"``. + #. The dimensions of ``lpcosts`` are the decision-variable dimensions. They are + renamed ``"d(xxx)"``. + #. ``set(b.dims).intersection(lpcosts.xxx.dims)`` are diagonal + in constraint dimensions and decision variables dimension, with ``xxx`` the + capacity or the production + #. ``set(constraint.xxx.dims) - set(lpcosts.xxx.dims) - set(b.dims)`` are reduced by + summation, with ``xxx`` the capacity or the production + #. ``set(lpcosts.xxx.dims) - set(constraint.xxx.dims) - set(b.dims)`` are added for + expansion, with ``xxx`` the capacity or the production + + See :py:func:`muse.lp_adapter.lp_constraint_matrix` for a more detailed explanation + of the transformations applied here. + """ constraint = constraint.copy(deep=False) @@ -295,7 +341,27 @@ def lp_constraint(constraint, lpcosts: xr.Dataset) -> xr.Dataset: def lp_constraint_matrix( b: xr.DataArray, constraint: xr.DataArray, lpcosts: xr.DataArray ): - """Transforms one constraint block into an lp matrix.""" + """Transforms one constraint block into an lp matrix. + + The goal is to create from ``lpcosts``, ``constraint``, and ``b`` a 2d-matrix of + constraints vs decision variables. + + #. The dimensions of ``b`` are the constraint dimensions. They are renamed + ``"c(xxx)"``. + #. The dimensions of ``lpcosts`` are the decision-variable dimensions. They are + renamed ``"d(xxx)"``. + #. ``set(b.dims).intersection(lpcosts.dims)`` are diagonal + in constraint dimensions and decision variables dimension + #. ``set(constraint.dims) - set(lpcosts.dims) - set(b.dims)`` are reduced by + summation + #. ``set(lpcosts.dims) - set(constraint.dims) - set(b.dims)`` are added for + expansion + #. ``set(b.dims) - set(constraint.dims) - set(lpcosts.dims)`` are added for + expansion. Such dimensions only make sense if they consist of one point. + + The result is the constraint matrix, expanded, reduced and diagonalized for the + conditions above. + """ from functools import reduce from numpy import eye From fb378143e31d83187f1eb5fcd270c109ef60c47e Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Tue, 10 Jun 2025 11:39:24 +0100 Subject: [PATCH 32/33] Update tests --- src/muse/investments.py | 2 +- src/muse/lp_adapter.py | 8 +- tests/test_constraints.py | 349 +++++++++++++++++--------------------- 3 files changed, 162 insertions(+), 197 deletions(-) diff --git a/src/muse/investments.py b/src/muse/investments.py index 78be4490..8a514936 100644 --- a/src/muse/investments.py +++ b/src/muse/investments.py @@ -297,7 +297,7 @@ def scipy_match_demand( # Run scipy optimization with highs solver adapter = ScipyAdapter.from_muse_data( - costs=costs, + capacity_costs=costs, constraints=constraints, commodities=commodities, timeslice_level=timeslice_level, diff --git a/src/muse/lp_adapter.py b/src/muse/lp_adapter.py index 31f55e3c..5430f368 100644 --- a/src/muse/lp_adapter.py +++ b/src/muse/lp_adapter.py @@ -156,18 +156,14 @@ class ScipyAdapter: @classmethod def from_muse_data( cls, - costs: xr.DataArray, + capacity_costs: xr.DataArray, constraints: list, commodities: list[str], timeslice_level: str | None = None, ) -> ScipyAdapter: """Creates a ScipyAdapter from MUSE data structures.""" # Calculate costs for the linear problem - lpcosts = lp_costs( - capacity_costs=costs, - commodities=commodities, - timeslice_level=timeslice_level, - ) + lpcosts = lp_costs(capacity_costs, commodities, timeslice_level) # Create dataset from costs and constraints data = unified_dataset(lpcosts, *constraints) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index b3902c06..f22fee40 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -3,17 +3,23 @@ import numpy as np import pandas as pd import xarray as xr -from pytest import approx, fixture, mark, raises +from pytest import approx, fixture, raises from muse import examples from muse.constraints import ( demand, demand_limiting_capacity, - lp_costs, max_capacity_expansion, max_production, + minimum_service, +) +from muse.lp_adapter import ( + ScipyAdapter, + back_to_muse_quantity, + lp_constraint_matrix, + selected_quantity, + unified_dataset, ) -from muse.quantities import maximum_production from muse.timeslices import drop_timeslice from muse.utilities import broadcast_over_assets, interpolate_capacity, reduce_assets @@ -32,6 +38,8 @@ def model_data(): - demand: Demand for the investment year - search_space: Search space for the assets """ + from muse.quantities import maximum_production + # Load residential sector data residential = examples.sector("residential", model="medium") @@ -72,12 +80,15 @@ def model_data(): @fixture(params=["timeslice_as_list", "timeslice_as_multindex"]) def constraints(request, model_data): + """Set of default constraints for testing.""" constraints = { "max_production": max_production(**model_data), "demand": demand(**model_data), "max_capacity_expansion": max_capacity_expansion(**model_data), "demand_limiting_capacity": demand_limiting_capacity(**model_data), } + + # Testing two different ways of handling timeslices if request.param == "timeslice_as_multindex": constraints = {key: _as_list(cs) for key, cs in constraints.items()} return constraints @@ -95,28 +106,6 @@ def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Da return data -@fixture -def lp_inputs(model_data): - """Inputs to the lp adapter, in addition to the constraints.""" - # Make up capacity costs data - shape = model_data["search_space"].shape - costs = model_data["search_space"] * np.arange(np.prod(shape)).reshape(shape) - - # List of commodities - commodities = list(model_data["demand"].commodity.values) - - return { - "capacity_costs": costs, - "commodities": commodities, - } - - -@fixture -def lpcosts(lp_inputs): - """Benchmark lpcosts dataset to test against.""" - return lp_costs(**lp_inputs) - - def test_model_data(model_data): """Validating that the fixture data has appropriate dimensions.""" assert set(model_data["technologies"].dims) == {"technology", "commodity"} @@ -129,25 +118,6 @@ def test_model_data(model_data): } -@mark.usefixtures("save_registries") -def test_objective_registration(): - from muse.objectives import OBJECTIVES, register_objective - - @register_objective - def a_objective(*args, **kwargs): - pass - - assert "a_objective" in OBJECTIVES - assert OBJECTIVES["a_objective"] is a_objective - - @register_objective(name="something") - def b_objective(*args, **kwargs): - pass - - assert "something" in OBJECTIVES - assert OBJECTIVES["something"] is b_objective - - def test_constraints_dimensions(constraints): """Test dimensions of all constraint matrices.""" # Max production constraint @@ -176,10 +146,132 @@ def test_constraints_dimensions(constraints): assert set(constraints["max_capacity_expansion"].b.dims) == {"replacement"} +def test_max_capacity_expansion(constraints): + max_capacity_expansion = constraints["max_capacity_expansion"] + assert (max_capacity_expansion.capacity == 1).all() + assert max_capacity_expansion.production == 0 + assert max_capacity_expansion.b.dims == ("replacement",) + assert max_capacity_expansion.b.shape == (4,) + assert ( + max_capacity_expansion.replacement + == ["estove", "gasboiler", "gasstove", "heatpump"] + ).all() + + +def test_max_production(constraints): + assert (constraints["max_production"].capacity <= 0).all() + + +def test_demand_limiting_capacity(constraints): + demand_limiting_capacity = constraints["demand_limiting_capacity"] + max_production = constraints["max_production"] + demand_constraint = constraints["demand"] + + # Test capacity values + expected_capacity = ( + -max_production.capacity.max("timeslice").values + if "timeslice" in max_production.capacity.dims + else -max_production.capacity.values + ) + assert demand_limiting_capacity.capacity.values == approx(expected_capacity) + + # Test production and b values + assert demand_limiting_capacity.production == 0 + expected_b = ( + demand_constraint.b.max("timeslice").values + if "timeslice" in demand_constraint.b.dims + else demand_constraint.b.values + ) + assert demand_limiting_capacity.b.values == approx(expected_b) + + +def test_minimum_service(model_data): + # Test with no minimum service factor (default) + assert minimum_service(**model_data) is None + + # Test with minimum service factor + technologies = model_data["technologies"].copy() + technologies["minimum_service_factor"] = 0.4 * xr.ones_like( + technologies.technology, dtype=float + ) + min_service = minimum_service(**{**model_data, "technologies": technologies}) + assert isinstance(min_service, xr.Dataset) + + +def test_max_capacity_expansion_no_limits(model_data): + technologies = model_data["technologies"].drop_vars( + ["max_capacity_addition", "max_capacity_growth", "total_capacity_limit"] + ) + assert ( + max_capacity_expansion(**{**model_data, "technologies": technologies}) is None + ) + + +def test_max_capacity_expansion_infinite_limits(model_data): + technologies = model_data["technologies"].copy() + for limit in [ + "max_capacity_addition", + "max_capacity_growth", + "total_capacity_limit", + ]: + technologies[limit] = np.inf + with raises(ValueError): + max_capacity_expansion(**{**model_data, "technologies": technologies}) + + +def test_max_capacity_expansion_seed(model_data): + seed = 10 + technologies = model_data["technologies"].copy() + technologies["growth_seed"] = seed + + # Test different capacity scenarios + scenarios = [0, seed, 2 * seed] + results = [] + for cap in scenarios: + capacity = model_data["capacity"].copy() + capacity.sel(year=2020)[:] = cap + results.append( + max_capacity_expansion( + **{ + **model_data, + "technologies": technologies, + "capacity": capacity, + } + ) + ) + + # Zero capacity should match seed capacity + assert results[0].b.values == approx(results[1].b.values) + # Higher capacity should differ + assert results[0].b.values != approx(results[2].b.values) + + +@fixture +def lp_inputs(model_data): + """Inputs to the lp adapter, in addition to the constraints.""" + # Make up capacity costs data + shape = model_data["search_space"].shape + costs = model_data["search_space"] * np.arange(np.prod(shape)).reshape(shape) + + # List of commodities + commodities = list(model_data["demand"].commodity.values) + + return { + "capacity_costs": costs, + "commodities": commodities, + } + + +@fixture +def lpcosts(lp_inputs): + """Benchmark lpcosts dataset to test against.""" + from muse.lp_adapter import lp_costs + + return lp_costs(**lp_inputs) + + def test_lp_constraints_matrix_b_is_scalar(lpcosts, constraints): """B is a scalar - output should be equivalent to a single row matrix.""" - from muse.constraints import lp_constraint_matrix - constraint = constraints["max_production"] for attr in ["capacity", "production"]: @@ -195,8 +287,6 @@ def test_lp_constraints_matrix_b_is_scalar(lpcosts, constraints): def test_max_production_constraint_diagonal(lpcosts, constraints): """Test production side of max capacity production is diagonal.""" - from muse.constraints import lp_constraint_matrix - constraint = constraints["max_production"] # Test capacity constraints @@ -224,7 +314,7 @@ def test_max_production_constraint_diagonal(lpcosts, constraints): def test_lp_constraint(lpcosts, constraints): - from muse.constraints import lp_constraint + from muse.lp_adapter import lp_constraint constraint = constraints["max_production"] @@ -252,9 +342,7 @@ def test_lp_constraint(lpcosts, constraints): def test_to_scipy_adapter_maxprod(lp_inputs, lpcosts, constraints): """Test scipy adapter with max production constraint.""" - from muse.constraints import ScipyAdapter - - adapter = ScipyAdapter.factory( + adapter = ScipyAdapter.from_muse_data( **lp_inputs, constraints=[constraints["max_production"]], ) @@ -290,9 +378,7 @@ def test_to_scipy_adapter_maxprod(lp_inputs, lpcosts, constraints): def test_to_scipy_adapter_demand(lp_inputs, lpcosts, constraints): """Test scipy adapter with demand constraint.""" - from muse.constraints import ScipyAdapter - - adapter = ScipyAdapter.factory( + adapter = ScipyAdapter.from_muse_data( **lp_inputs, constraints=[constraints["demand"]], ) @@ -322,9 +408,7 @@ def test_to_scipy_adapter_demand(lp_inputs, lpcosts, constraints): def test_to_scipy_adapter_max_capacity_expansion(lp_inputs, lpcosts, constraints): """Test scipy adapter with max capacity expansion constraint.""" - from muse.constraints import ScipyAdapter - - adapter = ScipyAdapter.factory( + adapter = ScipyAdapter.from_muse_data( **lp_inputs, constraints=[constraints["max_capacity_expansion"]], ) @@ -353,9 +437,7 @@ def test_to_scipy_adapter_max_capacity_expansion(lp_inputs, lpcosts, constraints def test_scipy_adapter_no_constraint(lp_inputs, lpcosts): - from muse.constraints import ScipyAdapter - - adapter = ScipyAdapter.factory( + adapter = ScipyAdapter.from_muse_data( **lp_inputs, constraints=[], ) @@ -370,38 +452,34 @@ def test_scipy_adapter_no_constraint(lp_inputs, lpcosts): def test_back_to_muse_quantities(lpcosts): - from muse.constraints import ScipyAdapter - - data = ScipyAdapter._unified_dataset(lpcosts) + data = unified_dataset(lpcosts) # Test capacity - lpquantity = ScipyAdapter._selected_quantity(data, "capacity") + lpquantity = selected_quantity(data, "capacity") assert set(lpquantity.dims) == {"d(asset)", "d(replacement)"} - copy = ScipyAdapter._back_to_muse_quantity( + copy = back_to_muse_quantity( lpquantity.costs.values, xr.zeros_like(lpquantity.costs) ) assert (copy == lpcosts.capacity).all() # Test production - lpquantity = ScipyAdapter._selected_quantity(data, "production") + lpquantity = selected_quantity(data, "production") assert set(lpquantity.dims) == { "d(asset)", "d(replacement)", "d(timeslice)", "d(commodity)", } - copy = ScipyAdapter._back_to_muse_quantity( + copy = back_to_muse_quantity( lpquantity.costs.values, xr.zeros_like(lpquantity.costs) ) assert (copy == lpcosts.production).all() -def test_back_to_muse_all(lpcosts): - from muse.constraints import ScipyAdapter - - data = ScipyAdapter._unified_dataset(lpcosts) - lpcapacity = ScipyAdapter._selected_quantity(data, "capacity") - lpproduction = ScipyAdapter._selected_quantity(data, "production") +def test_back_to_muse_all(lpcosts, lp_inputs): + data = unified_dataset(lpcosts) + lpcapacity = selected_quantity(data, "capacity") + lpproduction = selected_quantity(data, "production") x = np.concatenate( ( @@ -414,18 +492,17 @@ def test_back_to_muse_all(lpcosts): ) ) - copy = ScipyAdapter._back_to_muse( - x, xr.zeros_like(lpcapacity.costs), xr.zeros_like(lpproduction.costs) - ) + adapter = ScipyAdapter.from_muse_data(**lp_inputs, constraints=[]) + copy = adapter.to_muse(x) assert copy.capacity.size + copy.production.size == x.size assert (copy.capacity == lpcosts.capacity).all() assert (copy.production == lpcosts.production).all() def test_scipy_adapter_standard_constraints(lp_inputs, constraints): - from muse.constraints import ScipyAdapter - - adapter = ScipyAdapter.factory(**lp_inputs, constraints=constraints.values()) + adapter = ScipyAdapter.from_muse_data( + **lp_inputs, constraints=list(constraints.values()) + ) n_constraints = adapter.b_ub.size n_decision_vars = adapter.c.size @@ -446,119 +523,11 @@ def test_scipy_solver(model_data, lp_inputs, constraints): from muse.investments import scipy_match_demand solution = scipy_match_demand( - capacity_costs=lp_inputs["capacity_costs"], + costs=lp_inputs["capacity_costs"], commodities=lp_inputs["commodities"], search_space=model_data["search_space"], technologies=model_data["technologies"], - constraints=constraints.values(), + constraints=list(constraints.values()), ) assert isinstance(solution, xr.DataArray) assert set(solution.dims) == {"asset", "replacement"} - - -def test_minimum_service(model_data): - from muse.constraints import minimum_service - - # Test with no minimum service factor (default) - assert minimum_service(**model_data) is None - - # Test with minimum service factor - technologies = model_data["technologies"].copy() - technologies["minimum_service_factor"] = 0.4 * xr.ones_like( - technologies.technology, dtype=float - ) - min_service = minimum_service(**{**model_data, "technologies": technologies}) - assert isinstance(min_service, xr.Dataset) - - -def test_max_capacity_expansion_properties(constraints): - max_capacity_expansion = constraints["max_capacity_expansion"] - assert (max_capacity_expansion.capacity == 1).all() - assert max_capacity_expansion.production == 0 - assert max_capacity_expansion.b.dims == ("replacement",) - assert max_capacity_expansion.b.shape == (4,) - assert ( - max_capacity_expansion.replacement - == ["estove", "gasboiler", "gasstove", "heatpump"] - ).all() - - -def test_max_capacity_expansion_no_limits(model_data): - from muse.constraints import max_capacity_expansion - - technologies = model_data["technologies"].drop_vars( - ["max_capacity_addition", "max_capacity_growth", "total_capacity_limit"] - ) - assert ( - max_capacity_expansion(**{**model_data, "technologies": technologies}) is None - ) - - -def test_max_capacity_expansion_seed(model_data): - from muse.constraints import max_capacity_expansion - - seed = 10 - technologies = model_data["technologies"].copy() - technologies["growth_seed"] = seed - - # Test different capacity scenarios - scenarios = [0, seed, 2 * seed] - results = [] - for cap in scenarios: - capacity = model_data["capacity"].copy() - capacity.sel(year=2020)[:] = cap - results.append( - max_capacity_expansion( - **{ - **model_data, - "technologies": technologies, - "capacity": capacity, - } - ) - ) - - # Zero capacity should match seed capacity - assert results[0].b.values == approx(results[1].b.values) - # Higher capacity should differ - assert results[0].b.values != approx(results[2].b.values) - - -def test_max_capacity_expansion_infinite_limits(model_data): - from muse.constraints import max_capacity_expansion - - technologies = model_data["technologies"].copy() - for limit in [ - "max_capacity_addition", - "max_capacity_growth", - "total_capacity_limit", - ]: - technologies[limit] = np.inf - with raises(ValueError): - max_capacity_expansion(**{**model_data, "technologies": technologies}) - - -def test_max_production(constraints): - assert (constraints["max_production"].capacity <= 0).all() - - -def test_demand_limiting_capacity(constraints): - demand_limiting_capacity = constraints["demand_limiting_capacity"] - max_production = constraints["max_production"] - demand_constraint = constraints["demand"] - - # Test capacity values - expected_capacity = ( - -max_production.capacity.max("timeslice").values - if "timeslice" in max_production.capacity.dims - else -max_production.capacity.values - ) - assert demand_limiting_capacity.capacity.values == approx(expected_capacity) - - # Test production and b values - assert demand_limiting_capacity.production == 0 - expected_b = ( - demand_constraint.b.max("timeslice").values - if "timeslice" in demand_constraint.b.dims - else demand_constraint.b.values - ) - assert demand_limiting_capacity.b.values == approx(expected_b) From 5bd271ef294e6f76d859e6a8a88f4fb35ee76aee Mon Sep 17 00:00:00 2001 From: Tom Bland Date: Tue, 10 Jun 2025 12:06:09 +0100 Subject: [PATCH 33/33] Add MSF to default constraints --- tests/test_constraints.py | 43 ++++++++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index f22fee40..42bfb934 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -47,7 +47,12 @@ def model_data(): technologies = residential.technologies.squeeze("region").sel(year=INVESTMENT_YEAR) assets = next(a.assets for a in residential.agents) - # Calculate capacity and market demand + # Add minimum service factor data to allow calculation of the constraint + technologies["minimum_service_factor"] = 0.1 * xr.ones_like( + technologies.technology, dtype=float + ) + + # Calculate capacity capacity = interpolate_capacity( reduce_assets(assets.capacity, coords=("technology", "region")), year=[CURRENT_YEAR, INVESTMENT_YEAR], @@ -60,6 +65,7 @@ def model_data(): ).sel(year=INVESTMENT_YEAR).groupby("technology").sum("asset").rename( technology="asset" ) + # Remove un-demanded commodities market_demand = market_demand.sel( commodity=(market_demand > 0).any(dim=["timeslice", "asset"]) @@ -80,12 +86,13 @@ def model_data(): @fixture(params=["timeslice_as_list", "timeslice_as_multindex"]) def constraints(request, model_data): - """Set of default constraints for testing.""" + """Default set of constraints for testing.""" constraints = { "max_production": max_production(**model_data), "demand": demand(**model_data), "max_capacity_expansion": max_capacity_expansion(**model_data), "demand_limiting_capacity": demand_limiting_capacity(**model_data), + "minimum_service": minimum_service(**model_data), } # Testing two different ways of handling timeslices @@ -107,7 +114,7 @@ def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Da def test_model_data(model_data): - """Validating that the fixture data has appropriate dimensions.""" + """Validating that the model data has appropriate dimensions.""" assert set(model_data["technologies"].dims) == {"technology", "commodity"} assert set(model_data["search_space"].dims) == {"asset", "replacement"} assert set(model_data["capacity"].dims) == {"asset", "year"} @@ -145,8 +152,14 @@ def test_constraints_dimensions(constraints): assert set(constraints["max_capacity_expansion"].production.dims) == set() assert set(constraints["max_capacity_expansion"].b.dims) == {"replacement"} + # Minimum service constraint + assert set(constraints["minimum_service"].capacity.dims) == max_prod_dims + assert set(constraints["minimum_service"].production.dims) == max_prod_dims + assert set(constraints["minimum_service"].b.dims) == max_prod_dims + def test_max_capacity_expansion(constraints): + """Checking basic properties of the max capacity expansion constraint.""" max_capacity_expansion = constraints["max_capacity_expansion"] assert (max_capacity_expansion.capacity == 1).all() assert max_capacity_expansion.production == 0 @@ -159,10 +172,12 @@ def test_max_capacity_expansion(constraints): def test_max_production(constraints): + """Checking basic properties of the max production constraint.""" assert (constraints["max_production"].capacity <= 0).all() def test_demand_limiting_capacity(constraints): + """Checking basic properties of the demand limiting capacity constraint.""" demand_limiting_capacity = constraints["demand_limiting_capacity"] max_production = constraints["max_production"] demand_constraint = constraints["demand"] @@ -185,20 +200,8 @@ def test_demand_limiting_capacity(constraints): assert demand_limiting_capacity.b.values == approx(expected_b) -def test_minimum_service(model_data): - # Test with no minimum service factor (default) - assert minimum_service(**model_data) is None - - # Test with minimum service factor - technologies = model_data["technologies"].copy() - technologies["minimum_service_factor"] = 0.4 * xr.ones_like( - technologies.technology, dtype=float - ) - min_service = minimum_service(**{**model_data, "technologies": technologies}) - assert isinstance(min_service, xr.Dataset) - - def test_max_capacity_expansion_no_limits(model_data): + """Checking that the constraint is None when no limits are set.""" technologies = model_data["technologies"].drop_vars( ["max_capacity_addition", "max_capacity_growth", "total_capacity_limit"] ) @@ -208,6 +211,7 @@ def test_max_capacity_expansion_no_limits(model_data): def test_max_capacity_expansion_infinite_limits(model_data): + """Checking that error is raised when infinite limits are set.""" technologies = model_data["technologies"].copy() for limit in [ "max_capacity_addition", @@ -220,6 +224,7 @@ def test_max_capacity_expansion_infinite_limits(model_data): def test_max_capacity_expansion_seed(model_data): + """Sanity checks for the seed parameter of the max capacity expansion constraint.""" seed = 10 technologies = model_data["technologies"].copy() technologies["growth_seed"] = seed @@ -246,6 +251,12 @@ def test_max_capacity_expansion_seed(model_data): assert results[0].b.values != approx(results[2].b.values) +def test_no_minimum_service(model_data): + """Checking that the constraint is None when no minimum service factor is set.""" + technologies = model_data["technologies"].drop_vars("minimum_service_factor") + assert minimum_service(**{**model_data, "technologies": technologies}) is None + + @fixture def lp_inputs(model_data): """Inputs to the lp adapter, in addition to the constraints."""