Feature: piecewise constraints

CHANGELOG.md

@@ -1,3 +1,10 @@
+## 0.7.0.dev4 (dev)
+
+### User-facing changes
+
+|new| Piecewise constraints added to the YAML math with its own unique syntax.
+These constraints will be added to the optimisation problem using Special Ordered Sets of Type 2 (SOS2) variables.
+
 ## 0.7.0.dev3 (2024-02-14)
 
 ### User-facing changes

docs/migrating.md

@@ -977,5 +977,7 @@ Now, all components of our internal math are defined in a readable YAML syntax t
 
 You can add your own math to update the pre-defined math and to represent the physical system in ways we do not cover in our base math, or to apply new modelling methods and problem types (e.g., pathway or stochastic optimisation)!
 
+When adding your own math, you can add [piecewise linear constraints](user_defined_math/components.md#piecewise-constraints), which is a new type of constraint compared to what could be defined in v0.6.
+
 !!! info "See also"
     Our [pre-defined](pre_defined_math/index.md) and [user-defined](user_defined_math/index.md) math documentation.

docs/user_defined_math/components.md

@@ -78,15 +78,56 @@ constraints:
 ```
 
 1. It needs a unique name (`set_storage_initial` in the above example).
-1. Ideally, it has a long-form `description` and a `unit` added.
-These are not required, but are useful metadata for later reference.
+1. Ideally, it has a long-form `description` added.
+This is not required, but is useful metadata for later reference.
 1. It can have a top-level `foreach` list and `where` string.
 Without a `foreach`, it becomes an un-indexed constraint.
 Without a `where` string, all valid members (according to the `definition_matrix`) based on `foreach` will be included in this constraint.
 1. It has [equations](syntax.md#equations) (and, optionally, [sub-expressions](syntax.md#sub-expressions) and [slices](syntax.md#slices)) with corresponding lists of `where`+`expression` dictionaries.
 The equation expressions _must_ have comparison operators.
 1. It can be deactivated so that it does not appear in the built optimisation problem by setting `active: false`.
 
+## Piecewise constraints
+
+If you have non-linear relationships between two decision variables, you may want to represent them as a [piecewise linear function](https://en.wikipedia.org/wiki/Piecewise_linear_function).
+The most common form of a piecewise function involves creating special ordered sets of type 2 (SOS2), set of binary variables that are linked together with specific constraints.
+
+Because the formulation of piecewise constraints is so specific, the math syntax differs from all other modelling components by having `x` and `y` attributes that need to be specified:
+
+```yaml
+piecewise_constraints:
+  sos2_piecewise_flow_out:
+    description: Set outflow to follow capacity according to a piecewise curve.
+    foreach: [nodes, techs, carriers]
+    where: piecewise_x AND piecewise_y
+    x_expression: flow_cap
+    x_values: piecewise_x
+    y_expression: flow_out
+    y_values: piecewise_y
+    active: true
+```
+
+1. It needs a unique name (`sos2_piecewise_flow_out` in the above example).
+1. Ideally, it has a long-form `description` added.
+This is not required, but is useful metadata for later reference.
+1. It can have a top-level `foreach` list and `where` string.
+Without a `foreach`, it becomes an un-indexed constraint.
+Without a `where` string, all valid members (according to the `definition_matrix`) based on `foreach` will be included in this constraint.
+1. It has `x` and `y` [expression strings](syntax.md#expression-strings) (`x_expression`, `y_expression`).
+1. It has `x` and `y` parameter references (`x_values`, `y_values`).
+This should be a string name referencing an input parameter that contains the `breakpoints` dimension.
+1. It can be deactivated so that it does not appear in the built optimisation problem by setting `active: false`.
+
+The component attributes combine to describe a piecewise curve that links the `x_expression` and `y_expression` according to their respective values in `x_values` and `y_values` at each breakpoint.
+
+!!! note
+    If the non-linear function you want to represent is convex, you may be able to avoid SOS2 variables, and instead represent it using [constraint components](#constraints).
+    You can find an example of this in our [piecewise linear costs example][piecewise-linear-costs].
+
+!!! warning
+    This approximation of a non-linear relationship may improve the representation of whatever real system you are modelling, but it will come at the cost of a more difficult model to solve.
+    Indeed, introducing piecewise constraints may mean your model can no longer reach a solution with the computational resources you have available.
+
 ## Objectives
 
 With your constrained decision variables and a global expression that binds these variables to costs, you need an objective to minimise/maximise. The default, pre-defined objective is `min_cost_optimisation` and looks as follows:

docs/user_defined_math/examples/sos2_piecewise_linear_costs.yaml

@@ -0,0 +1,32 @@
+# title: Piecewise linear costs - economies of scale
+#
+# description: |
+#     Add a piecewise cost function that reduces the incremental increase in investment costs with increasing technology rated capacity.
+#     This emulates "economies of scale", where the more of a technology there is deployed, the less expensive each additional investment in deployment.
+#
+#     New indexed parameters:
+#
+#     - `piecewise_cost_investment_x` (defining the new set `breakpoints`)
+#     - `piecewise_cost_investment_y` (defining the new set `breakpoints`)
+#
+# ---
+
+variables:
+  piecewise_cost_investment:
+    description: Investment cost that increases monotonically
+    foreach: [nodes, techs, carriers, costs]
+    where: any(piecewise_cost_investment_x, over=breakpoints) AND any(piecewise_cost_investment_y, over=breakpoints)
+    bounds:
+      min: 0
+      max: .inf
+
+piecewise_constraints:
+  sos2_piecewise_costs:
+    description: >
+      Set investment costs values along a piecewise curve using special ordered sets of type 2 (SOS2).
+    foreach: [nodes, techs, carriers, costs]
+    where: any(piecewise_cost_investment_x, over=breakpoints) AND any(piecewise_cost_investment_y, over=breakpoints)
+    x_expression: flow_cap
+    x_values: piecewise_cost_investment_x
+    y_expression: piecewise_cost_investment
+    y_values: piecewise_cost_investment_y

pyproject.toml

@@ -33,6 +33,9 @@ directory = "reports/coverage"
 [tool.coverage.xml]
 output = "reports/coverage/coverage.xml"
 
+[tool.coverage.report]
+exclude_lines = ["if TYPE_CHECKING:"]
+
 [tool.black]
 line-length = 88
 skip-magic-trailing-comma = true

src/calliope/backend/backend_model.py

@@ -48,15 +48,20 @@
 
 T = TypeVar("T")
 _COMPONENTS_T = Literal[
-    "variables", "constraints", "objectives", "parameters", "global_expressions"
+    "parameters",
+    "variables",
+    "global_expressions",
+    "constraints",
+    "piecewise_constraints",
+    "objectives",
 ]
 
 LOGGER = logging.getLogger(__name__)
 
 
 class BackendModelGenerator(ABC):
     _VALID_COMPONENTS: tuple[_COMPONENTS_T, ...] = typing.get_args(_COMPONENTS_T)
-    _COMPONENT_ATTR_METADATA = ["description", "unit", "default"]
+    _COMPONENT_ATTR_METADATA = ["description", "unit", "default", "math_repr"]
 
     _PARAM_DESCRIPTIONS = extract_from_schema(MODEL_SCHEMA, "description")
     _PARAM_UNITS = extract_from_schema(MODEL_SCHEMA, "x-unit")
@@ -119,6 +124,21 @@ def add_constraint(
                 Constraint configuration dictionary, ready to be parsed and then evaluated.
         """
 
+    @abstractmethod
+    def add_piecewise_constraint(
+        self, name: str, constraint_dict: parsing.UnparsedPiecewiseConstraintDict
+    ) -> None:
+        """
+        Add piecewise constraint equation to backend model in-place.
+        Resulting backend dataset entries will be piecewise constraint objects.
+
+        Args:
+            name (str):
+                Name of the piecewise constraint
+            constraint_dict (parsing.UnparsedPiecewiseConstraintDict):
+                Piecewise constraint configuration dictionary, ready to be parsed and then evaluated.
+        """
+
     @abstractmethod
     def add_global_expression(
         self, name: str, expression_dict: parsing.UnparsedExpressionDict
@@ -216,12 +236,13 @@ def _build(self) -> None:
             "variables",
             "global_expressions",
             "constraints",
+            "piecewise_constraints",
             "objectives",
         ]:
             component = components.removesuffix("s")
-            for name in self.inputs.math[components]:
+            for name, dict_ in self.inputs.math[components].items():
                 start = time.time()
-                getattr(self, f"add_{component}")(name)
+                getattr(self, f"add_{component}")(name, dict_)
                 end = time.time() - start
                 LOGGER.debug(
                     f"Optimisation Model | {components}:{name} | Built in {end:.4f}s"
@@ -252,10 +273,14 @@ def _add_run_mode_math(self) -> None:
     def _add_component(
         self,
         name: str,
-        component_dict: Optional[Tp],
+        component_dict: Tp,
         component_setter: Callable,
         component_type: Literal[
-            "variables", "global_expressions", "constraints", "objectives"
+            "variables",
+            "global_expressions",
+            "constraints",
+            "piecewise_constraints",
+            "objectives",
         ],
         break_early: bool = True,
     ) -> Optional[parsing.ParsedBackendComponent]:
@@ -279,9 +304,6 @@ def _add_component(
         """
         references: set[str] = set()
 
-        if component_dict is None:
-            component_dict = self.inputs.math[component_type][name]
-
         if break_early and not component_dict.get("active", True):
             self.log(
                 component_type, name, "Component deactivated and therefore not built."
@@ -454,16 +476,29 @@ def _add_to_dataset(
             }
         )
         self._dataset[name] = da
-
         if references is not None:
-            for reference in references:
-                try:
-                    self._dataset[reference].attrs["references"].add(name)
-                except KeyError:
-                    continue
+            self._update_references(name, references)
+
+    def _update_references(self, name: str, references: set):
+        """Update reference lists in dataset objects.
+
+        Args:
+            name (str): Name to update in reference lists.
+            references (set): Names of dataset objects whose reference lists will be updated with `name`.
+        """
+        for reference in references:
+            try:
+                self._dataset[reference].attrs["references"].add(name)
+            except KeyError:
+                continue
 
     def _apply_func(
-        self, func: Callable, *args, output_core_dims: tuple = ((),), **kwargs
+        self,
+        func: Callable,
+        *args,
+        input_core_dims: Optional[list] = None,
+        output_core_dims: tuple = ((),),
+        **kwargs,
     ) -> xr.DataArray:
         """
         Apply a function to every element of an arbitrary number of xarray DataArrays.
@@ -476,10 +511,14 @@ def _apply_func(
             args (xr.DataArray):
                 xarray DataArrays which will be broadcast together and then iterated over
                 to apply the function.
-            output_core_dims (tuple):
+            input_core_dims (Optional[tuple], optional):
+                Additional dimensions which `xr.apply_ufunc` won't broadcast on applying `func`.
+                This is directly passed to `xr.apply_ufunc`; see their documentation for more details.
+                Defaults to None.
+            output_core_dims (tuple, optional):
                 Additional dimensions which are expected to be passed back from `xr.apply_ufunc` after applying `func`.
                 This is directly passed to `xr.apply_ufunc`; see their documentation for more details.
-                Defaults to ((), )
+                Defaults to ((), ).
             kwargs (dict[str, Any]):
                 Additional keyword arguments to pass to `func`.
 
@@ -495,6 +534,7 @@ def _apply_func(
             dask="parallelized",
             output_dtypes=[np.dtype("O")],
             output_core_dims=output_core_dims,
+            input_core_dims=input_core_dims,
         )
 
     def _raise_error_on_preexistence(self, key: str, obj_type: _COMPONENTS_T):
@@ -530,6 +570,11 @@ def constraints(self):
         "Slice of backend dataset to show only built constraints"
         return self._dataset.filter_by_attrs(obj_type="constraints")
 
+    @property
+    def piecewise_constraints(self):
+        "Slice of backend dataset to show only built constraints"
+        return self._dataset.filter_by_attrs(obj_type="piecewise_constraints")
+
     @property
     def variables(self):
         "Slice of backend dataset to show only built variables"
@@ -621,6 +666,21 @@ def get_constraint(
                 Otherwise, a xr.Dataset will be given, indexed over the same dimensions as the xr.DataArray, with variables for the constraint body, and upper (`ub`) and lower (`lb`) bounds.
         """
 
+    def get_piecewise_constraint(self, name: str) -> xr.DataArray:
+        """Get piecewise constraint data as an array of backend interface objects.
+        Can be used to inspect and debug built piecewise constraints.
+
+        Unlike other optimisation problem components, piecewise constraints can only be inspected as backend interface objects.
+        This is because each element is a collection of variables, parameters, constraints, and expressions.
+
+        Args:
+            name (str): Name of piecewise constraint, as given in YAML piecewise constraint key.
+
+        Returns:
+            xr.DataArray: Piecewise constraint array.
+        """
+        return self._get_component(name, "piecewise_constraints")
+
     @abstractmethod
     def get_variable(self, name: str, as_backend_objs: bool = True) -> xr.DataArray:
         """Extract decision variable array from backend dataset
@@ -851,22 +911,23 @@ def _rebuild_references(self, references: set[str]) -> None:
         Args:
             references (set[str]): names of optimisation problem components.
         """
-        ordered_components = [
-            "parameters",
-            "variables",
-            "global_expressions",
-            "constraints",
-            "objectives",
-        ]
-        for component in ordered_components:
+        for component in self._VALID_COMPONENTS:
             refs = [
                 ref
                 for ref in references
                 if self._dataset[ref].attrs["obj_type"] == component
             ]
             for ref in refs:
                 self.delete_component(ref, component)
-                getattr(self, "add_" + component.removesuffix("s"))(name=ref)
+                dict_ = self.inputs.attrs["math"][component][ref]
+                getattr(self, "add_" + component.removesuffix("s"))(ref, dict_)
+
+    def _get_component(self, name: str, component_group: str) -> xr.DataArray:
+        component = getattr(self, component_group).get(name, None)
+        if component is None:
+            pretty_group_name = component_group.removesuffix("s").replace("_", " ")
+            raise KeyError(f"Unknown {pretty_group_name}: {name}")
+        return component
 
 
 class ShadowPrices:

src/calliope/backend/expression_parser.py

@@ -736,15 +736,12 @@ def as_math_string(self) -> str:
         evaluated = self.as_array()
         self.eval_attrs["references"].add(self.name)
 
-        if evaluated.shape:
-            dims = rf"_\text{{{','.join(str(i).removesuffix('s') for i in evaluated.dims)}}}"
+        if evaluated.attrs["obj_type"] != "string":
+            data_var_string = evaluated.attrs["math_repr"]
         else:
-            dims = ""
-        if evaluated.attrs["obj_type"] in ["global_expressions", "variables"]:
-            formatted_name = rf"\textbf{{{self.name}}}"
-        elif evaluated.attrs["obj_type"] == "parameters":
-            formatted_name = rf"\textit{{{self.name}}}"
-        return formatted_name + dims
+            data_var_string = rf"\text{{{self.name}}}"
+
+        return data_var_string
 
     def as_array(self) -> xr.DataArray:
         backend_interface = self.eval_attrs["backend_interface"]