Add Variable.create_fixed, support expressions for variable bounds, i…

…mprove set printing, denser .lp files (#53)

* Minor improvements

* Add support for expression upper bound and strict __setattr__

* Add support for variables fixed to an expression

* Improve set printing

* Rename variable method

* Improved file writing

* Reformat with black

.gitignore

@@ -9,4 +9,5 @@ tmp
 .coverage
 .cache
 docs/reference
-dist
+dist
+*.log

README.md

@@ -17,3 +17,9 @@ Contributions are welcome! See [`CONTRIBUTE.md`](./CONTRIBUTE.md).
 ## Acknowledgments
 
 Martin Staadecker first created this library while working for [Bravos Power](https://www.bravospower.com/) The library takes inspiration from Linopy and Pyomo, two prior libraries for optimization for which we are thankful.
+
+## Troubleshooting Common Errors
+
+### `datatypes of join keys don't match`
+
+Often, this error indicates that two dataframes in your inputs representing the same dimension have different datatypes (e.g. 16bit integer and 64bit integer). This is not allowed and you should ensure for the same dimensions, datatypes are identical.

pyproject.toml

@@ -16,7 +16,7 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
     "Natural Language :: English",
 ]
-dependencies = ["polars==0.20.13", "numpy", "pyarrow", "pandas"]
+dependencies = ["polars", "numpy", "pyarrow", "pandas", "tqdm"]
 
 [project.optional-dependencies]
 dev = [

src/pyoframe/constants.py

@@ -50,6 +50,8 @@ class Config(metaclass=_ConfigMeta):
     disable_unmatched_checks: bool = False
     print_float_precision: Optional[int] = 5
     print_uses_variable_names: bool = True
+    # Number of elements to show when printing a set to the console (additional elements are replaced with ...)
+    print_max_set_elements: int = 50
 
     @classmethod
     def reset_defaults(cls):

src/pyoframe/constraints.py

@@ -35,6 +35,7 @@
     get_obj_repr,
     parse_inputs_as_iterable,
     unwrap_single_values,
+    dataframe_to_tupled_list,
 )
 
 from pyoframe.model_element import (
@@ -248,7 +249,9 @@ def __repr__(self):
         return (
             get_obj_repr(self, ("name",), size=self.data.height, dimensions=self.shape)
             + "\n"
-            + self.to_expr().to_str(max_line_len=80, max_rows=10)
+            + dataframe_to_tupled_list(
+                self.data, num_max_elements=Config.print_max_set_elements
+            )
         )
 
     @staticmethod
@@ -375,18 +378,18 @@ def map(self, mapping_set: SetTypes, drop_shared_dims: bool = True):
 
         >>> import polars as pl
         >>> from pyoframe import Variable, Model
-        >>> pop_data = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "population": [10, 2, 8]}).to_expr()
+        >>> pop_data = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "year": [2024, 2024, 2024], "population": [10, 2, 8]}).to_expr()
         >>> cities_and_countries = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "country": ["Canada", "Canada", "USA"]})
         >>> pop_data.map(cities_and_countries)
-        <Expression size=2 dimensions={'country': 2} terms=2>
-        [Canada]: 12
-        [USA]: 8
+        <Expression size=2 dimensions={'year': 1, 'country': 2} terms=2>
+        [2024,Canada]: 12
+        [2024,USA]: 8
 
         >>> pop_data.map(cities_and_countries, drop_shared_dims=False)
-        <Expression size=3 dimensions={'city': 3, 'country': 2} terms=3>
-        [Toronto,Canada]: 10
-        [Vancouver,Canada]: 2
-        [Boston,USA]: 8
+        <Expression size=3 dimensions={'city': 3, 'year': 1, 'country': 2} terms=3>
+        [Toronto,2024,Canada]: 10
+        [Vancouver,2024,Canada]: 2
+        [Boston,2024,USA]: 8
         """
         mapping_set = Set(mapping_set)
 

src/pyoframe/io.py

@@ -6,6 +6,7 @@
 from tempfile import NamedTemporaryFile
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterable, Optional, TypeVar, Union
+from tqdm import tqdm
 
 from pyoframe.constants import CONST_TERM, VAR_KEY
 from pyoframe.constraints import Constraint
@@ -39,7 +40,9 @@ def objective_to_file(m: "Model", f: TextIOWrapper, var_map):
 
 
 def constraints_to_file(m: "Model", f: TextIOWrapper, var_map, const_map):
-    for constraint in create_section(m.constraints, f, "s.t."):
+    for constraint in create_section(
+        tqdm(m.constraints, desc="Writing constraints to file"), f, "s.t."
+    ):
         f.writelines(constraint.to_str(var_map=var_map, const_map=const_map) + "\n")
 
 
@@ -55,17 +58,25 @@ def bounds_to_file(m: "Model", f, var_map):
         )
         f.write(f"{var_map.apply(const_term_df).item()} = 1\n")
 
-    for variable in m.variables:
-        lb = f"{variable.lb:.12g}"
-        ub = f"{variable.ub:.12g}"
+    for variable in tqdm(m.variables, desc="Writing bounds to file"):
+        terms = []
+
+        if variable.lb != 0:
+            terms.append(pl.lit(f"{variable.lb:.12g} <= "))
+
+        terms.append(VAR_KEY)
+
+        if variable.ub != float("inf"):
+            terms.append(pl.lit(f" <= {variable.ub:.12g}"))
+
+        terms.append(pl.lit("\n"))
+
+        if len(terms) < 3:
+            continue
 
         df = (
             var_map.apply(variable.data, to_col=None)
-            .select(
-                pl.concat_str(
-                    pl.lit(f"{lb} <= "), VAR_KEY, pl.lit(f" <= {ub}\n")
-                ).str.concat("")
-            )
+            .select(pl.concat_str(terms).str.concat(""))
             .item()
         )
 
@@ -76,7 +87,9 @@ def binaries_to_file(m: "Model", f, var_map: Mapper):
     """
     Write out binaries of a model to a lp file.
     """
-    for variable in create_section(m.binary_variables, f, "binary"):
+    for variable in create_section(
+        tqdm(m.binary_variables, "Writing binary variables to file"), f, "binary"
+    ):
         lines = (
             var_map.apply(variable.data, to_col=None)
             .select(pl.col(VAR_KEY).str.concat("\n"))
@@ -89,7 +102,9 @@ def integers_to_file(m: "Model", f, var_map: Mapper):
     """
     Write out integers of a model to a lp file.
     """
-    for variable in create_section(m.integer_variables, f, "general"):
+    for variable in create_section(
+        tqdm(m.integer_variables, "Writing integer variables to file"), f, "general"
+    ):
         lines = (
             var_map.apply(variable.data, to_col=None)
             .select(pl.col(VAR_KEY).str.concat("\n"))
@@ -155,6 +170,6 @@ def to_file(
         bounds_to_file(m, f, var_map)
         binaries_to_file(m, f, var_map)
         integers_to_file(m, f, var_map)
-        f.write("end\n")
+        f.write("\nend\n")
 
     return file_path

src/pyoframe/model.py

@@ -1,5 +1,5 @@
 from typing import Any, Iterable, List, Optional
-from pyoframe.constants import ObjSense, VType, Config, Result
+from pyoframe.constants import ObjSense, VType, Config, Result, PyoframeError
 from pyoframe.constraints import SupportsMath
 from pyoframe.io_mappers import NamedVariableMapper, IOMappers
 from pyoframe.model_element import ModelElement
@@ -9,13 +9,29 @@
 from pyoframe.variables import Variable
 from pyoframe.io import to_file
 from pyoframe.solvers import solve, Solver
+import polars as pl
+import pandas as pd
 
 
 class Model(AttrContainerMixin):
     """
     Represents a mathematical optimization model. Add variables, constraints, and an objective to the model by setting attributes.
     """
 
+    _reserved_attributes = [
+        "_variables",
+        "_constraints",
+        "_objective",
+        "var_map",
+        "io_mappers",
+        "name",
+        "solver",
+        "solver_model",
+        "params",
+        "result",
+        "attr",
+    ]
+
     def __init__(self, name=None, **kwargs):
         super().__init__(**kwargs)
         self._variables: List[Variable] = []
@@ -62,6 +78,13 @@ def minimize(self):
         return self.objective
 
     def __setattr__(self, __name: str, __value: Any) -> None:
+        if __name not in Model._reserved_attributes and not isinstance(
+            __value, (ModelElement, pl.DataFrame, pd.DataFrame)
+        ):
+            raise PyoframeError(
+                f"Cannot set attribute '{__name}' on the model because it isn't of type ModelElement (e.g. Variable, Constraint, ...)"
+            )
+
         if __name in ("maximize", "minimize"):
             assert isinstance(
                 __value, SupportsMath
@@ -71,24 +94,25 @@ def __setattr__(self, __name: str, __value: Any) -> None:
             self._objective._model = self
             return
 
-        if isinstance(__value, ModelElement) and not __name.startswith("_"):
+        if (
+            isinstance(__value, ModelElement)
+            and __name not in Model._reserved_attributes
+        ):
             assert not hasattr(
                 self, __name
             ), f"Cannot create {__name} since it was already created."
 
-            __value.name = __name
-            __value._model = self
+            __value.on_add_to_model(self, __name)
 
             if isinstance(__value, Objective):
                 assert self.objective is None, "Cannot create more than one objective."
                 self._objective = __value
-            if isinstance(__value, Variable):
+            elif isinstance(__value, Variable):
                 self._variables.append(__value)
                 if self.var_map is not None:
                     self.var_map.add(__value)
             elif isinstance(__value, Constraint):
                 self._constraints.append(__value)
-
         return super().__setattr__(__name, __value)
 
     def __repr__(self) -> str:

src/pyoframe/model_element.py

@@ -39,6 +39,10 @@ def __init__(self, data: pl.DataFrame, **kwargs) -> None:
         self.name = None
         super().__init__(**kwargs)
 
+    def on_add_to_model(self, model: "Model", name: str):
+        self.name = name
+        self._model = model
+
     @property
     def data(self) -> pl.DataFrame:
         return self._data
@@ -106,7 +110,11 @@ def _support_polars_method(method_name: str):
     """
 
     def method(self: "SupportPolarsMethodMixin", *args, **kwargs):
-        return self._new(getattr(self.data, method_name)(*args, **kwargs))
+        result_from_polars = getattr(self.data, method_name)(*args, **kwargs)
+        if isinstance(result_from_polars, pl.DataFrame):
+            return self._new(result_from_polars)
+        else:
+            return result_from_polars
 
     return method
 
@@ -115,6 +123,7 @@ class SupportPolarsMethodMixin:
     rename = _support_polars_method("rename")
     with_columns = _support_polars_method("with_columns")
     filter = _support_polars_method("filter")
+    estimated_size = _support_polars_method("estimated_size")
 
     @abstractmethod
     def _new(self, data: pl.DataFrame):

src/pyoframe/util.py

@@ -229,3 +229,36 @@ def wrapper(*args, **kwargs):
         return result
 
     return wrapper
+
+
+def dataframe_to_tupled_list(
+    df: pl.DataFrame, num_max_elements: Optional[int] = None
+) -> str:
+    """
+    Converts a dataframe into a list of tuples. Used to print a Set to the console. See examples for behaviour.
+
+    Examples:
+        >>> df = pl.DataFrame({"x": [1, 2, 3, 4, 5]})
+        >>> dataframe_to_tupled_list(df)
+        '[1, 2, 3, 4, 5]'
+        >>> dataframe_to_tupled_list(df, 3)
+        '[1, 2, 3, ...]'
+
+        >>> df = pl.DataFrame({"x": [1, 2, 3, 4, 5], "y": [2, 3, 4, 5, 6]})
+        >>> dataframe_to_tupled_list(df, 3)
+        '[(1, 2), (2, 3), (3, 4), ...]'
+    """
+    elipse = False
+    if num_max_elements is not None:
+        if len(df) > num_max_elements:
+            elipse = True
+            df = df.head(num_max_elements)
+
+    res = (row for row in df.iter_rows())
+    if len(df.columns) == 1:
+        res = (row[0] for row in res)
+
+    res = str(list(res))
+    if elipse:
+        res = res[:-1] + ", ...]"
+    return res

src/pyoframe/variables.py

@@ -3,7 +3,7 @@
 """
 
 from __future__ import annotations
-from typing import Any, Iterable
+from typing import Iterable, TYPE_CHECKING
 
 import polars as pl
 
@@ -17,11 +17,14 @@
     VType,
     VTypeValue,
 )
-from pyoframe.constraints import Expression
+from pyoframe.constraints import Expression, SupportsToExpr
 from pyoframe.constraints import SetTypes
 from pyoframe.util import get_obj_repr, unwrap_single_values
 from pyoframe.model_element import CountableModelElement, SupportPolarsMethodMixin
 
+if TYPE_CHECKING:
+    from pyoframe.model import Model
+
 
 class Variable(CountableModelElement, SupportsMath, SupportPolarsMethodMixin):
     """
@@ -65,21 +68,44 @@ class Variable(CountableModelElement, SupportsMath, SupportPolarsMethodMixin):
     def __init__(
         self,
         *indexing_sets: SetTypes | Iterable[SetTypes],
-        lb: float = float("-inf"),
-        ub: float = float("inf"),
+        lb: float | int | SupportsToExpr = float("-inf"),
+        ub: float | int | SupportsToExpr = float("inf"),
         vtype: VType | VTypeValue = VType.CONTINUOUS,
     ):
         data = Set(*indexing_sets).data if len(indexing_sets) > 0 else pl.DataFrame()
         super().__init__(data)
 
         self.vtype: VType = VType(vtype)
+        self._fixed_to = None
 
         # Tightening the bounds is not strictly necessary, but it adds clarity
         if self.vtype == VType.BINARY:
             lb, ub = 0, 1
 
-        self.lb = lb
-        self.ub = ub
+        if isinstance(lb, (float, int)):
+            self.lb, self.lb_constraint = lb, None
+        else:
+            self.lb, self.lb_constraint = float("-inf"), lb <= self
+
+        if isinstance(ub, (float, int)):
+            self.ub, self.ub_constraint = ub, None
+        else:
+            self.ub, self.ub_constraint = float("inf"), self <= ub
+
+    def on_add_to_model(self, model: "Model", name: str):
+        super().on_add_to_model(model, name)
+        if self.lb_constraint is not None:
+            setattr(model, f"{name}_lb", self.lb_constraint)
+        if self.ub_constraint is not None:
+            setattr(model, f"{name}_ub", self.ub_constraint)
+        if self._fixed_to is not None:
+            setattr(model, f"{name}_fixed", self == self._fixed_to)
+
+    @classmethod
+    def create_fixed(cls, expr: SupportsToExpr):
+        v = Variable(expr)
+        v._fixed_to = expr
+        return v
 
     @classmethod
     def get_id_column_name(cls):