feat: add sos constraints

Author: coroa

linopy/io.py

@@ -23,6 +23,7 @@
 from tqdm import tqdm
 
 from linopy import solvers
+from linopy.common import to_polars
 from linopy.constants import CONCAT_DIM
 from linopy.objective import Objective
 
@@ -327,6 +328,66 @@ def integers_to_file(
             formatted.write_csv(f, **kwargs)
 
 
+def sos_to_file(
+    m: Model,
+    f: BufferedWriter,
+    progress: bool = False,
+    slice_size: int = 2_000_000,
+    explicit_coordinate_names: bool = False,
+) -> None:
+    """
+    Write out integers of a model to a lp file.
+    """
+    names = m.variables.sos
+    if not len(list(names)):
+        return
+
+    print_variable, _ = get_printers(
+        m, explicit_coordinate_names=explicit_coordinate_names
+    )
+
+    f.write(b"\n\nsos\n\n")
+    if progress:
+        names = tqdm(
+            list(names),
+            desc="Writing sos constraints.",
+            colour=TQDM_COLOR,
+        )
+
+    for name in names:
+        var = m.variables[name]
+        sos_type = var.attrs["sos_type"]
+        sos_dim = var.attrs["sos_dim"]
+
+        other_dims = tuple([dim for dim in var.labels.dims if dim != sos_dim])
+        for var_slice in var.iterate_slices(slice_size, other_dims):
+            ds = var_slice.labels.to_dataset()
+            ds["sos_labels"] = ds["labels"].isel({sos_dim: 0})
+            ds["weights"] = ds.coords[sos_dim]
+            df = to_polars(ds)
+
+            df = df.group_by("sos_labels").agg(
+                pl.concat_str(
+                    *print_variable(pl.col("labels")), pl.lit(":"), pl.col("weights")
+                )
+                .str.join(" ")
+                .alias("var_weights")
+            )
+
+            columns = [
+                pl.lit("s"),
+                pl.col("sos_labels"),
+                pl.lit(f": S{sos_type} :: "),
+                pl.col("var_weights"),
+            ]
+
+            kwargs: Any = dict(
+                separator=" ", null_value="", quote_style="never", include_header=False
+            )
+            formatted = df.select(pl.concat_str(columns, ignore_nulls=True))
+            formatted.write_csv(f, **kwargs)
+
+
 def constraints_to_file(
     m: Model,
     f: BufferedWriter,
@@ -464,6 +525,13 @@ def to_lp_file(
             slice_size=slice_size,
             explicit_coordinate_names=explicit_coordinate_names,
         )
+        sos_to_file(
+            m,
+            f=f,
+            progress=progress,
+            slice_size=slice_size,
+            explicit_coordinate_names=explicit_coordinate_names,
+        )
         f.write(b"end\n")
 
         logger.info(f" Writing time: {round(time.time() - start, 2)}s")

linopy/model.py

@@ -12,7 +12,7 @@
 from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from tempfile import NamedTemporaryFile, gettempdir
-from typing import Any, overload
+from typing import Any, Literal, overload
 
 import numpy as np
 import pandas as pd
@@ -551,6 +551,39 @@ def add_variables(
         self.variables.add(variable)
         return variable
 
+    def add_sos_constraints(
+        self,
+        variable: Variable,
+        sos_type: Literal[1, 2],
+        sos_dim: str,
+    ):
+        """
+        Add an sos1 or sos2 constraint for one dimension of a variable
+
+        The dimension values are used as SOS.
+
+        Parameters
+        ----------
+        variable : Variable
+        sos_type : {1, 2}
+            Type of SOS
+        sos_dim : str
+            Which dimension of variable to add SOS constraint to
+        """
+        if sos_type not in (1, 2):
+            raise ValueError(f"sos_type must be 1 or 2, got {sos_type}")
+        if sos_dim not in variable.dims:
+            raise ValueError(f"sos_dim must name a variable dimension, got {sos_dim}")
+
+        if "sos_type" in variable.attrs or "sos_dim" in variable.attrs:
+            sos_type = variable.attrs.get("sos_type")
+            sos_dim = variable.attrs.get("sos_dim")
+            raise ValueError(
+                "variable already has an sos{sos_type} constraint on {sos_dim}"
+            )
+
+        variable.attrs.update(sos_type=sos_type, sos_dim=sos_dim)
+
     def add_constraints(
         self,
         lhs: VariableLike

linopy/variables.py

@@ -193,6 +193,14 @@ def __init__(
         if "label_range" not in data.attrs:
             data.assign_attrs(label_range=(data.labels.min(), data.labels.max()))
 
+        if "sos_type" in data.attrs or "sos_dim" in data.attrs:
+            if (sos_type := data.attrs.get("sos_type")) not in (1, 2):
+                raise ValueError(f"sos_type must be 1 or 2, got {sos_type}")
+            if (sos_dim := data.attrs.get("sos_dim")) not in data.dims:
+                raise ValueError(
+                    f"sos_dim must name a variable dimension, got {sos_dim}"
+                )
+
         self._data = data
         self._model = model
 
@@ -323,6 +331,8 @@ def __repr__(self) -> str:
         dim_names = self.coord_names
         dim_sizes = list(self.sizes.values())
         masked_entries = (~self.mask).sum().values
+        sos_type = self.attrs.get("sos_type")
+        sos_dim = self.attrs.get("sos_dim")
         lines = []
 
         if dims:
@@ -344,9 +354,11 @@ def __repr__(self) -> str:
 
             shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes))
             mask_str = f" - {masked_entries} masked entries" if masked_entries else ""
+            sos_str = f" - sos{sos_type} on {sos_dim}" if sos_type and sos_dim else ""
             lines.insert(
                 0,
-                f"Variable ({shape_str}){mask_str}\n{'-' * (len(shape_str) + len(mask_str) + 11)}",
+                f"Variable ({shape_str}){mask_str}{sos_str}\n"
+                f"{'-' * (len(shape_str) + len(mask_str) + len(sos_str) + 11)}",
             )
         else:
             lines.append(
@@ -1362,6 +1374,21 @@ def continuous(self) -> Variables:
             self.model,
         )
 
+    @property
+    def sos(self) -> Variables:
+        """
+        Get all variables involved in an sos constraint.
+        """
+        return self.__class__(
+            {
+                name: self.data[name]
+                for name in self
+                if self[name].attrs.get("sos_dim")
+                and self[name].attrs.get("sos_type") in (1, 2)
+            },
+            self.model,
+        )
+
     @property
     def solution(self) -> Dataset:
         """