fix: extra output comparison

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

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)

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
 

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)