pyoframe 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

pyoframe/io.py

@@ -6,16 +6,17 @@ from io import TextIOWrapper
 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 VAR_KEY, Config
-from pyoframe.constraints import Constraint
-from pyoframe.variables import Variable
+from pyoframe.constants import CONST_TERM, VAR_KEY, ObjSense
+from pyoframe.core import Constraint, Variable
 from pyoframe.io_mappers import (
     Base62ConstMapper,
     Base62VarMapper,
     IOMappers,
     Mapper,
     NamedMapper,
+    NamedVariableMapper,
 )
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -28,63 +29,88 @@ def objective_to_file(m: "Model", f: TextIOWrapper, var_map):
     """
     Write out the objective of a model to a lp file.
     """
-    assert m.objective is not None, "No objective set."
-
-    f.write(f"{m.objective.sense.value}\n\nobj:\n\n")
-    result = m.objective.to_str(var_map=var_map, include_prefix=False)
-    f.writelines(result)
+    if m.objective is None:
+        return
+    objective_sense = "minimize" if m.sense == ObjSense.MIN else "maximize"
+    f.write(f"{objective_sense}\n\nobj:\n\n")
+    result = m.objective.to_str(
+        var_map=var_map, include_prefix=False, include_const_variable=True
+    )
+    f.write(result)
 
 
 def constraints_to_file(m: "Model", f: TextIOWrapper, var_map, const_map):
-    for constraint in create_section(m.constraints, f, "s.t."):
-        f.writelines(constraint.to_str(var_map=var_map, const_map=const_map) + "\n")
+    for constraint in create_section(
+        tqdm(m.constraints, desc="Writing constraints to file"), f, "s.t."
+    ):
+        f.write(constraint.to_str(var_map=var_map, const_map=const_map) + "\n")
 
 
 def bounds_to_file(m: "Model", f, var_map):
     """
     Write out variables of a model to a lp file.
     """
-    for variable in create_section(m.variables, f, "bounds"):
-        lb = f"{variable.lb:.12g}"
-        ub = f"{variable.ub:.12g}"
+    if (m.objective is not None and m.objective.has_constant) or len(m.variables) != 0:
+        f.write("\n\nbounds\n\n")
+    if m.objective is not None and m.objective.has_constant:
+        const_term_df = pl.DataFrame(
+            {VAR_KEY: [CONST_TERM]}, schema={VAR_KEY: pl.UInt32}
+        )
+        f.write(f"{var_map.apply(const_term_df).item()} = 1\n")
+
+    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()
         )
 
-        f.writelines(df)
+        f.write(df)
 
 
 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"))
             .item()
         )
-        f.writelines(lines + "\n")
+        f.write(lines + "\n")
 
 
 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"))
             .item()
         )
-        f.writelines(lines + "\n")
+        f.write(lines + "\n")
 
 
 T = TypeVar("T")
@@ -103,7 +129,7 @@ def get_var_map(m: "Model", use_var_names):
     if use_var_names:
         if m.var_map is not None:
             return m.var_map
-        var_map = NamedMapper(Variable)
+        var_map = NamedVariableMapper(Variable)
     else:
         var_map = Base62VarMapper(Variable)
 
@@ -144,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

pyoframe/io_mappers.py

@@ -9,14 +9,14 @@ from abc import ABC, abstractmethod
 
 from typing import TYPE_CHECKING, Optional, Type, Union
 import polars as pl
-from pyoframe.constants import NAME_COL
 from pyoframe.util import concat_dimensions
+from pyoframe.constants import CONST_TERM
 
 
 if TYPE_CHECKING:  # pragma: no cover
     from pyoframe.model import Variable
-    from pyoframe.constraints import Constraint
-    from pyoframe.util import IdCounterMixin
+    from pyoframe.core import Constraint
+    from pyoframe.util import CountableModelElement
 
 
 @dataclass
@@ -26,38 +26,47 @@ class IOMappers:
 
 
 class Mapper(ABC):
-    def __init__(self, cls: Type["IdCounterMixin"]) -> None:
+
+    NAME_COL = "__name"
+
+    def __init__(self, cls: Type["CountableModelElement"]) -> None:
         self._ID_COL = cls.get_id_column_name()
         self.mapping_registry = pl.DataFrame(
-            {self._ID_COL: [], NAME_COL: []},
-            schema={self._ID_COL: pl.UInt32, NAME_COL: pl.String},
+            {self._ID_COL: [], Mapper.NAME_COL: []},
+            schema={self._ID_COL: pl.UInt32, Mapper.NAME_COL: pl.String},
         )
 
     def add(self, element: Union["Variable", "Constraint"]) -> None:
-        self.mapping_registry = pl.concat(
-            [self.mapping_registry, self._element_to_map(element)]
-        )
+        self._extend_registry(self._element_to_map(element))
+
+    def _extend_registry(self, df: pl.DataFrame) -> None:
+        self.mapping_registry = pl.concat([self.mapping_registry, df])
 
     @abstractmethod
-    def _element_to_map(self, element: "IdCounterMixin") -> pl.DataFrame: ...
+    def _element_to_map(self, element: "CountableModelElement") -> pl.DataFrame: ...
 
     def apply(
         self,
         df: pl.DataFrame,
-        to_col: Optional[str],
+        to_col: Optional[str] = None,
     ) -> pl.DataFrame:
+        if df.height == 0:
+            return df
         result = df.join(
             self.mapping_registry, on=self._ID_COL, how="left", validate="m:1"
         )
         if to_col is None:
             result = result.drop(self._ID_COL)
             to_col = self._ID_COL
-        return result.rename({NAME_COL: to_col})
+        return result.rename({Mapper.NAME_COL: to_col})
 
     def undo(self, df: pl.DataFrame) -> pl.DataFrame:
+        if df.height == 0:
+            return df
+        df = df.rename({self._ID_COL: Mapper.NAME_COL})
         return df.join(
-            self.mapping_registry, on=NAME_COL, how="left", validate="m:1"
-        ).drop(NAME_COL)
+            self.mapping_registry, on=Mapper.NAME_COL, how="left", validate="m:1"
+        ).drop(Mapper.NAME_COL)
 
 
 class NamedMapper(Mapper):
@@ -68,7 +77,7 @@ class NamedMapper(Mapper):
 
         >>> import polars as pl
         >>> import pyoframe as pf
-        >>> m = pf.Model()
+        >>> m = pf.Model("min")
         >>> m.foo = pf.Variable(pl.DataFrame({"t": range(4)}))
         >>> pf.sum(m.foo)
         <Expression size=1 dimensions={} terms=4>
@@ -81,7 +90,20 @@ class NamedMapper(Mapper):
             element_name is not None
         ), "Element must have a name to be used in a named mapping."
         return concat_dimensions(
-            element.ids, keep_dims=False, prefix=element_name, to_col=NAME_COL
+            element.ids, keep_dims=False, prefix=element_name, to_col=Mapper.NAME_COL
+        )
+
+
+class NamedVariableMapper(NamedMapper):
+    CONST_TERM_NAME = "_ONE"
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._extend_registry(
+            pl.DataFrame(
+                {self._ID_COL: [CONST_TERM], self.NAME_COL: [self.CONST_TERM_NAME]},
+                schema={self._ID_COL: pl.UInt32, self.NAME_COL: pl.String},
+            )
         )
 
 
@@ -168,7 +190,7 @@ class Base62Mapper(Mapper, ABC):
         )
 
     def _element_to_map(self, element) -> pl.DataFrame:
-        return self.apply(element.ids.select(self._ID_COL), to_col=NAME_COL)
+        return self.apply(element.ids.select(self._ID_COL), to_col=Mapper.NAME_COL)
 
 
 class Base62VarMapper(Base62Mapper):
@@ -177,7 +199,7 @@ class Base62VarMapper(Base62Mapper):
         >>> import polars as pl
         >>> from pyoframe import Model, Variable
         >>> from pyoframe.constants import VAR_KEY
-        >>> m = Model()
+        >>> m = Model("min")
         >>> m.x = Variable(pl.DataFrame({"t": range(1,63)}))
         >>> (m.x.filter(t=11)+1).to_str()
         '[11]: 1  + x[11]'
@@ -194,6 +216,15 @@ class Base62VarMapper(Base62Mapper):
         └───────────────┘
     """
 
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        df = pl.DataFrame(
+            {self._ID_COL: [CONST_TERM]},
+            schema={self._ID_COL: pl.UInt32},
+        )
+        df = self.apply(df, to_col=Mapper.NAME_COL)
+        self._extend_registry(df)
+
     @property
     def _prefix(self) -> "str":
         return "x"

pyoframe/model.py

@@ -1,29 +1,59 @@
-from typing import Any, Iterable, List, Optional
-from pyoframe.constants import ObjSense, VType, Config
-from pyoframe.constraints import SupportsMath
-from pyoframe.io_mappers import NamedMapper, IOMappers
-from pyoframe.model_element import ModelElement
-from pyoframe.constraints import Constraint
+from typing import Any, Iterable, List, Optional, Union
+from pyoframe.constants import (
+    ObjSense,
+    VType,
+    Config,
+    Result,
+    PyoframeError,
+    ObjSenseValue,
+)
+from pyoframe.io_mappers import NamedVariableMapper, IOMappers
+from pyoframe.model_element import ModelElement, ModelElementWithId
+from pyoframe.core import Constraint
 from pyoframe.objective import Objective
-from pyoframe.variables import Variable
+from pyoframe.user_defined import Container, AttrContainerMixin
+from pyoframe.core import Variable
 from pyoframe.io import to_file
-from pyoframe.solvers import solve
+from pyoframe.solvers import solve, Solver
+import polars as pl
+import pandas as pd
 
 
-class Model:
+class Model(AttrContainerMixin):
     """
     Represents a mathematical optimization model. Add variables, constraints, and an objective to the model by setting attributes.
     """
 
-    def __init__(self, name=None):
+    _reserved_attributes = [
+        "_variables",
+        "_constraints",
+        "_objective",
+        "var_map",
+        "io_mappers",
+        "name",
+        "solver",
+        "solver_model",
+        "params",
+        "result",
+        "attr",
+        "sense",
+    ]
+
+    def __init__(self, min_or_max: Union[ObjSense, ObjSenseValue], name=None, **kwargs):
+        super().__init__(**kwargs)
         self._variables: List[Variable] = []
         self._constraints: List[Constraint] = []
+        self.sense = ObjSense(min_or_max)
         self._objective: Optional[Objective] = None
         self.var_map = (
-            NamedMapper(Variable) if Config.print_uses_variable_names else None
+            NamedVariableMapper(Variable) if Config.print_uses_variable_names else None
         )
         self.io_mappers: Optional[IOMappers] = None
         self.name = name
+        self.solver: Optional[Solver] = None
+        self.solver_model: Optional[Any] = None
+        self.params = Container()
+        self.result: Optional[Result] = None
 
     @property
     def variables(self) -> List[Variable]:
@@ -45,44 +75,37 @@ class Model:
     def objective(self):
         return self._objective
 
-    @property
-    def maximize(self):
-        assert self.objective is not None and self.objective.sense == ObjSense.MAX
-        return self.objective
-
-    @property
-    def minimize(self):
-        assert self.objective is not None and self.objective.sense == ObjSense.MIN
-        return self.objective
-
     def __setattr__(self, __name: str, __value: Any) -> None:
-        if __name in ("maximize", "minimize"):
-            assert isinstance(
-                __value, SupportsMath
-            ), f"Setting {__name} on the model requires an objective expression."
-            self._objective = Objective(__value, sense=__name)
-            self._objective.name = __name
-            self._objective._model = self
-            return
-
-        if isinstance(__value, ModelElement) and not __name.startswith("_"):
-            assert not hasattr(
-                self, __name
-            ), f"Cannot create {__name} since it was already created."
-
-            __value.name = __name
-            __value._model = self
-
-            if isinstance(__value, Objective):
-                assert self.objective is None, "Cannot create more than one objective."
-                self._objective = __value
+        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 (
+            isinstance(__value, ModelElement)
+            and __name not in Model._reserved_attributes
+        ):
+            if __name == "objective":
+                __value = Objective(__value)
+
+            if isinstance(__value, ModelElementWithId):
+                assert not hasattr(
+                    self, __name
+                ), f"Cannot create {__name} since it was already created."
+
+            __value.on_add_to_model(self, __name)
+
             if 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)
-
+            elif isinstance(__value, Objective):
+                self._objective = __value
+                return
         return super().__setattr__(__name, __value)
 
     def __repr__(self) -> str:

pyoframe/model_element.py

@@ -1,27 +1,18 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from collections import defaultdict
+from typing import Any, Dict, List, Optional
 import polars as pl
 from typing import TYPE_CHECKING
 
 from pyoframe.constants import COEF_KEY, RESERVED_COL_KEYS, VAR_KEY
 from pyoframe._arithmetic import _get_dimensions
+from pyoframe.user_defined import AttrContainerMixin
 
 if TYPE_CHECKING:  # pragma: no cover
     from pyoframe.model import Model
 
 
-def _pass_polars_method(method_name: str):
-    """
-    Wrapper to add a method to ModelElement that simply calls the underlying Polars method on the data attribute.
-    """
-
-    def method(self, *args, **kwargs):
-        return self._new(getattr(self.data, method_name)(*args, **kwargs))
-
-    return method
-
-
 class ModelElement(ABC):
     def __init__(self, data: pl.DataFrame, **kwargs) -> None:
         # Sanity checks, no duplicate column names
@@ -48,17 +39,25 @@ class ModelElement(ABC):
         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
 
+    @property
+    def friendly_name(self) -> str:
+        return self.name if self.name is not None else "unnamed"
+
     @property
     def dimensions(self) -> Optional[List[str]]:
         """
         The names of the data's dimensions.
 
         Examples:
-            >>> from pyoframe.variables import Variable
+            >>> from pyoframe.core import Variable
             >>> # A variable with no dimensions
             >>> Variable().dimensions
 
@@ -85,7 +84,7 @@ class ModelElement(ABC):
         The number of indices in each dimension.
 
         Examples:
-            >>> from pyoframe.variables import Variable
+            >>> from pyoframe.core import Variable
             >>> # A variable with no dimensions
             >>> Variable().shape
             {}
@@ -104,13 +103,120 @@ class ModelElement(ABC):
             return 1
         return self.data.select(dims).n_unique()
 
+
+def _support_polars_method(method_name: str):
+    """
+    Wrapper to add a method to ModelElement that simply calls the underlying Polars method on the data attribute.
+    """
+
+    def method(self: "SupportPolarsMethodMixin", *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
+
+
+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):
         """
         Used to create a new instance of the same class with the given data (for e.g. on .rename(), .with_columns(), etc.).
         """
-        raise NotImplementedError
 
-    rename = _pass_polars_method("rename")
-    with_columns = _pass_polars_method("with_columns")
-    filter = _pass_polars_method("filter")
+
+class ModelElementWithId(ModelElement, AttrContainerMixin):
+    """
+    Provides a method that assigns a unique ID to each row in a DataFrame.
+    IDs start at 1 and go up consecutively. No zero ID is assigned since it is reserved for the constant variable term.
+    IDs are only unique for the subclass since different subclasses have different counters.
+    """
+
+    # Keys are the subclass names and values are the next unasigned ID.
+    _id_counters: Dict[str, int] = defaultdict(lambda: 1)
+
+    @classmethod
+    def reset_counters(cls):
+        """
+        Resets all the ID counters.
+        This function is called before every unit test to reset the code state.
+        """
+        cls._id_counters = defaultdict(lambda: 1)
+
+    def __init__(self, data: pl.DataFrame, **kwargs) -> None:
+        super().__init__(data, **kwargs)
+        self._data = self._assign_ids(self.data)
+
+    @classmethod
+    def _assign_ids(cls, df: pl.DataFrame) -> pl.DataFrame:
+        """
+        Adds the column `to_column` to the DataFrame `df` with the next batch
+        of unique consecutive IDs.
+        """
+        cls_name = cls.__name__
+        cur_count = cls._id_counters[cls_name]
+        id_col_name = cls.get_id_column_name()
+
+        if df.height == 0:
+            df = df.with_columns(pl.lit(cur_count).alias(id_col_name))
+        else:
+            df = df.with_columns(
+                pl.int_range(cur_count, cur_count + pl.len()).alias(id_col_name)
+            )
+        df = df.with_columns(pl.col(id_col_name).cast(pl.UInt32))
+        cls._id_counters[cls_name] += df.height
+        return df
+
+    @classmethod
+    @abstractmethod
+    def get_id_column_name(cls) -> str:
+        """
+        Returns the name of the column containing the IDs.
+        """
+
+    @property
+    def ids(self) -> pl.DataFrame:
+        return self.data.select(self.dimensions_unsafe + [self.get_id_column_name()])
+
+    def _extend_dataframe_by_id(self, addition: pl.DataFrame):
+        cols = addition.columns
+        assert len(cols) == 2
+        id_col = self.get_id_column_name()
+        assert id_col in cols
+        cols.remove(id_col)
+        new_col = cols[0]
+
+        original = self.data
+
+        if new_col in original.columns:
+            original = original.drop(new_col)
+        self._data = original.join(addition, on=id_col, how="left", validate="1:1")
+
+    def _preprocess_attr(self, name: str, value: Any) -> Any:
+        dims = self.dimensions
+        ids = self.ids
+        id_col = self.get_id_column_name()
+
+        if isinstance(value, pl.DataFrame):
+            if value.shape == (1, 1):
+                value = value.item()
+            else:
+                assert (
+                    dims is not None
+                ), "Attribute must be a scalar since there are no dimensions"
+                result = value.join(ids, on=dims, validate="1:1", how="inner").drop(
+                    dims
+                )
+                assert len(result.columns) == 2, "Attribute has too many columns"
+                value_col = [c for c in result.columns if c != id_col][0]
+                return result.rename({value_col: name})
+
+        assert ids.height == 1, "Attribute is a scalar but there are multiple IDs."
+        return pl.DataFrame({name: [value], id_col: ids.get_column(id_col)})

pyoframe/monkey_patch.py

@@ -1,7 +1,7 @@
 import pandas as pd
 import polars as pl
-from pyoframe.constraints import SupportsMath
-from pyoframe.constraints import Expression
+from pyoframe.core import SupportsMath
+from pyoframe.core import Expression
 from functools import wraps
 
 from pyoframe.constants import COEF_KEY, CONST_TERM, VAR_KEY

pyoframe/objective.py

@@ -1,28 +1,24 @@
-from typing import Optional, Union
-from pyoframe.constants import ObjSense, ObjSenseValue
-from pyoframe.constraints import SupportsMath, Expression
+from typing import Optional
+from pyoframe.constants import COEF_KEY
+from pyoframe.core import SupportsMath, Expression
 
 
 class Objective(Expression):
     r"""
     Examples:
         >>> from pyoframe import Variable, Model, sum
-        >>> m = Model()
+        >>> m = Model("max")
         >>> m.a = Variable()
         >>> m.b = Variable({"dim1": [1, 2, 3]})
-        >>> m.maximize = m.a + sum("dim1", m.b)
-        >>> m.maximize
+        >>> m.objective = m.a + sum("dim1", m.b)
+        >>> m.objective
         <Objective size=1 dimensions={} terms=4>
-        maximize: a + b[1] + b[2] + b[3]
+        objective: a + b[1] + b[2] + b[3]
     """
 
-    def __init__(
-        self, expr: SupportsMath, sense: Union[ObjSense, ObjSenseValue]
-    ) -> None:
-        self.sense = ObjSense(sense)
-
+    def __init__(self, expr: SupportsMath) -> None:
         expr = expr.to_expr()
-        super().__init__(expr.to_expr().data)
+        super().__init__(expr.data)
         self._model = expr._model
         assert (
             self.dimensions is None
@@ -40,3 +36,10 @@ class Objective(Expression):
     @value.setter
     def value(self, value):
         self._value = value
+
+    @property
+    def has_constant(self):
+        constant_terms = self.constant_terms
+        if len(constant_terms) == 0:
+            return False
+        return constant_terms.get_column(COEF_KEY).item() != 0