diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index b901c7e6..5916016b 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -2,6 +2,19 @@ All notable changes to GemsPy are documented here. +## [Unreleased] + +### Fixed +- Comparison operators (`>=`, `<=`, `=`) in extra-output expressions no longer raise + `NotImplementedError` at post-solve evaluation; they now evaluate to a 0/1 indicator + (e.g. `dual(balance) >= unsupplied_cost - 5`). +- Extra-output `minimum()`/`maximum()` no longer break when several models coexist + in the same problem: solution and constraint-dual arrays are now filtered back + down to each model's own components before evaluation, instead of keeping the + NaN-padded entries introduced by linopy's outer join across models. + +--- + ## [0.1.2] - 2026-06-11 ### Added diff --git a/src/gems/simulation/extra_output.py b/src/gems/simulation/extra_output.py index bd5f1d78..6044d05a 100644 --- a/src/gems/simulation/extra_output.py +++ b/src/gems/simulation/extra_output.py @@ -26,12 +26,19 @@ """ from dataclasses import dataclass, field -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, cast import numpy as np import xarray as xr -from gems.expression.expression import DualNode, ReducedCostNode, VariableNode +from gems.expression.expression import ( + Comparator, + ComparisonNode, + DualNode, + ReducedCostNode, + VariableNode, +) +from gems.expression.visitor import visit from gems.model.port import PortFieldId from gems.simulation.vectorized_builder import VectorizedBuilderBase from gems.study.system import Component @@ -146,3 +153,22 @@ def reduced_cost(self, node: ReducedCostNode) -> xr.DataArray: f"{self.model_id!r}." ) return self.var_reduced_cost_arrays[key] + + def comparison(self, node: ComparisonNode) -> xr.DataArray: + """Evaluate a comparison post-solve as a float indicator DataArray. + + Both operands resolve to concrete ``xr.DataArray`` values (solved + variables, duals, parameters, ...), so the comparison is evaluated + element-wise and returned as ``1.0`` where it holds and ``0.0`` + otherwise. This is only meaningful post-solve; the pre-solve builder + keeps rejecting comparisons (they are constraints, not values). + """ + left = cast(xr.DataArray, visit(node.left, self)) + right = cast(xr.DataArray, visit(node.right, self)) + if node.comparator == Comparator.LESS_THAN: # "<=" + indicator = left <= right + elif node.comparator == Comparator.GREATER_THAN: # ">=" + indicator = left >= right + else: # Comparator.EQUAL -> "=" + indicator = left == right + return cast(xr.DataArray, indicator).astype(float) diff --git a/src/gems/simulation/simulation_table.py b/src/gems/simulation/simulation_table.py index e662079a..9f8abbea 100644 --- a/src/gems/simulation/simulation_table.py +++ b/src/gems/simulation/simulation_table.py @@ -271,7 +271,11 @@ def _collect_extra_outputs( if solution is not None: for (mk, vname), lv in problem._linopy_vars.items(): if lv.name in solution: - var_solution_arrays[(mk, vname)] = solution[lv.name] + sol_da = solution[lv.name] + if "component" in sol_da.dims: + own_components = list(lv.coords["component"].values) + sol_da = sol_da.sel(component=own_components) + var_solution_arrays[(mk, vname)] = sol_da constraint_dual_arrays = self._collect_constraint_duals(problem) var_reduced_cost_arrays = self._collect_reduced_costs(problem) @@ -339,9 +343,10 @@ def _collect_constraint_duals( """Return constraint shadow prices keyed by (model_key, constraint_name).""" dual_dataset = problem.linopy_model.dual result: Dict[Tuple[str, str], xr.DataArray] = {} - for mk in problem.study.model_components: + for mk, components in problem.study.model_components.items(): model = problem.study.models[mk] prefix = mk.replace("-", "_") + own_components = [c.id for c in components] all_constraints = {**model.constraints, **model.binding_constraints} for cname in all_constraints: safe = cname.replace(" ", "_").replace("-", "_") @@ -355,6 +360,8 @@ def _collect_constraint_duals( dual_val = dual_val + dual_dataset[lb_name] # type: ignore[operator] if ub_name in dual_dataset: dual_val = dual_val + dual_dataset[ub_name] # type: ignore[operator] + if "component" in dual_val.dims: + dual_val = dual_val.sel(component=own_components) result[(mk, cname)] = dual_val return result diff --git a/tests/unittests/simulation/test_simulation_table_extra_outputs.py b/tests/unittests/simulation/test_simulation_table_extra_outputs.py index 7e71d980..8dda87f2 100644 --- a/tests/unittests/simulation/test_simulation_table_extra_outputs.py +++ b/tests/unittests/simulation/test_simulation_table_extra_outputs.py @@ -162,3 +162,97 @@ def test_extra_output_abs_round_on_variable() -> None: ) assert abs_shift == pytest.approx(2.3) assert rounded == pytest.approx(3.0) + + +def test_extra_output_min_on_variable() -> None: + """ + min(variable, parameter) is allowed in extra outputs and evaluated + element-wise post-solve. + """ + from gems.expression import param, var + from gems.expression.expression import literal, minimum + from gems.model.model import model + from gems.model.parameter import float_parameter + from gems.model.variable import float_variable + from gems.simulation import TimeBlock, build_problem + from gems.study import ConstantData, DataBase, Study, System, create_component + + SIMPLE_MODEL = model( + id="SIMPLE_MIN", + parameters=[float_parameter("cap")], + variables=[float_variable("a", lower_bound=literal(5), upper_bound=literal(5))], + extra_outputs={"capped": minimum(var("a"), param("cap"))}, + ) + + database = DataBase() + comp = create_component(model=SIMPLE_MODEL, id="comp_1") + database.add_data("comp_1", "cap", ConstantData(3.0)) + + OTHER_MODEL = model( + id="OTHER_MODEL", + parameters=[float_parameter("p")], + variables=[float_variable("b", lower_bound=literal(1), upper_bound=literal(1))], + ) + other_comp = create_component(model=OTHER_MODEL, id="comp_2") + database.add_data("comp_2", "p", ConstantData(0.0)) + + system = System("test_min_extra") + system.add_component(comp) + system.add_component(other_comp) + + problem = build_problem( + Study(system, database), TimeBlock(1, [0]), scenario_ids=list(range(1)) + ) + problem.solve(solver_name="highs") + + df = SimulationTableBuilder().build(problem) + capped = ( + df.component("comp_1").output("capped").value(time_index=0, scenario_index=0) + ) + # min(a=5, cap=3) == 3 + assert capped == pytest.approx(3.0) + + +def test_extra_output_comparison() -> None: + """ + Comparison operators (>=, <=) are allowed in extra outputs and evaluated + post-solve as a float indicator: 1.0 where the condition holds, 0.0 + otherwise. Previously these raised NotImplementedError. + """ + from gems.expression import param, var + from gems.expression.expression import Comparator, ComparisonNode, literal + from gems.model.model import model + from gems.model.parameter import float_parameter + from gems.model.variable import float_variable + from gems.simulation import TimeBlock, build_problem + from gems.study import ConstantData, DataBase, Study, System, create_component + + SIMPLE_MODEL = model( + id="SIMPLE_CMP", + parameters=[float_parameter("threshold")], + variables=[float_variable("a", lower_bound=literal(5), upper_bound=literal(5))], + extra_outputs={ + # a (=5) >= threshold (=3) -> 1.0 + "ge": ComparisonNode(var("a"), param("threshold"), Comparator.GREATER_THAN), + # a (=5) <= threshold (=3) -> 0.0 + "le": ComparisonNode(var("a"), param("threshold"), Comparator.LESS_THAN), + }, + ) + + database = DataBase() + comp = create_component(model=SIMPLE_MODEL, id="comp_1") + database.add_data("comp_1", "threshold", ConstantData(3.0)) + + system = System("test_comparison_extra") + system.add_component(comp) + + problem = build_problem( + Study(system, database), TimeBlock(1, [0]), scenario_ids=list(range(1)) + ) + problem.solve(solver_name="highs") + + df = SimulationTableBuilder().build(problem) + ge = df.component("comp_1").output("ge").value(time_index=0, scenario_index=0) + le = df.component("comp_1").output("le").value(time_index=0, scenario_index=0) + assert ge == pytest.approx(1.0) + assert le == pytest.approx(0.0)