pyoframe 0.0.4__py3-none-any.whl

pyoframe/io.py

@@ -0,0 +1,149 @@
+"""
+Module containing all import/export functionalities.
+"""
+
+from io import TextIOWrapper
+from tempfile import NamedTemporaryFile
+from pathlib import Path
+from typing import TYPE_CHECKING, Iterable, Optional, TypeVar, Union
+
+from pyoframe.constants import VAR_KEY, Config
+from pyoframe.constraints import Constraint
+from pyoframe.variables import Variable
+from pyoframe.io_mappers import (
+    Base62ConstMapper,
+    Base62VarMapper,
+    IOMappers,
+    Mapper,
+    NamedMapper,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pyoframe.model import Model
+
+import polars as pl
+
+
+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)
+
+
+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")
+
+
+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}"
+
+        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("")
+            )
+            .item()
+        )
+
+        f.writelines(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"):
+        lines = (
+            var_map.apply(variable.data, to_col=None)
+            .select(pl.col(VAR_KEY).str.concat("\n"))
+            .item()
+        )
+        f.writelines(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"):
+        lines = (
+            var_map.apply(variable.data, to_col=None)
+            .select(pl.col(VAR_KEY).str.concat("\n"))
+            .item()
+        )
+        f.writelines(lines + "\n")
+
+
+T = TypeVar("T")
+
+
+def create_section(iterable: Iterable[T], f, section_header) -> Iterable[T]:
+    wrote = False
+    for item in iterable:
+        if not wrote:
+            f.write(f"\n\n{section_header}\n\n")
+            wrote = True
+        yield item
+
+
+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)
+    else:
+        var_map = Base62VarMapper(Variable)
+
+    for v in m.variables:
+        var_map.add(v)
+    return var_map
+
+
+def to_file(
+    m: "Model", file_path: Optional[Union[str, Path]], use_var_names=False
+) -> Path:
+    """
+    Write out a model to a lp file.
+    """
+    if file_path is None:
+        with NamedTemporaryFile(
+            prefix="pyoframe-problem-", suffix=".lp", mode="w", delete=False
+        ) as f:
+            file_path = f.name
+
+    file_path = Path(file_path)
+    assert file_path.suffix == ".lp", f"File format `{file_path.suffix}` not supported."
+
+    if file_path.exists():
+        file_path.unlink()
+
+    const_map = (
+        NamedMapper(Constraint) if use_var_names else Base62ConstMapper(Constraint)
+    )
+    for c in m.constraints:
+        const_map.add(c)
+    var_map = get_var_map(m, use_var_names)
+    m.io_mappers = IOMappers(var_map, const_map)
+
+    with open(file_path, mode="w") as f:
+        objective_to_file(m, f, var_map)
+        constraints_to_file(m, f, var_map, const_map)
+        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")
+
+    return file_path

pyoframe/io_mappers.py

@@ -0,0 +1,206 @@
+"""
+Defines various methods for mapping a variable or constraint to its string representation.
+"""
+
+from dataclasses import dataclass
+import math
+import string
+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
+
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pyoframe.model import Variable
+    from pyoframe.constraints import Constraint
+    from pyoframe.util import IdCounterMixin
+
+
+@dataclass
+class IOMappers:
+    var_map: "Mapper"
+    const_map: "Mapper"
+
+
+class Mapper(ABC):
+    def __init__(self, cls: Type["IdCounterMixin"]) -> 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},
+        )
+
+    def add(self, element: Union["Variable", "Constraint"]) -> None:
+        self.mapping_registry = pl.concat(
+            [self.mapping_registry, self._element_to_map(element)]
+        )
+
+    @abstractmethod
+    def _element_to_map(self, element: "IdCounterMixin") -> pl.DataFrame: ...
+
+    def apply(
+        self,
+        df: pl.DataFrame,
+        to_col: Optional[str],
+    ) -> pl.DataFrame:
+        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})
+
+    def undo(self, df: pl.DataFrame) -> pl.DataFrame:
+        return df.join(
+            self.mapping_registry, on=NAME_COL, how="left", validate="m:1"
+        ).drop(NAME_COL)
+
+
+class NamedMapper(Mapper):
+    """
+    Maps constraints or variables to a string representation using the object's name and dimensions.
+
+    Examples:
+
+        >>> import polars as pl
+        >>> import pyoframe as pf
+        >>> m = pf.Model()
+        >>> m.foo = pf.Variable(pl.DataFrame({"t": range(4)}))
+        >>> pf.sum(m.foo)
+        <Expression size=1 dimensions={} terms=4>
+        foo[0] + foo[1] + foo[2] + foo[3]
+    """
+
+    def _element_to_map(self, element) -> pl.DataFrame:
+        element_name = element.name  # type: ignore
+        assert (
+            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
+        )
+
+
+class Base62Mapper(Mapper, ABC):
+    # Mapping between a base 62 character and its integer value
+    _CHAR_TABLE = pl.DataFrame(
+        {"char": list(string.digits + string.ascii_letters)},
+    ).with_columns(pl.int_range(pl.len()).cast(pl.UInt32).alias("code"))
+
+    _BASE = _CHAR_TABLE.height  # _BASE = 62
+    _ZERO = _CHAR_TABLE.filter(pl.col("code") == 0).select("char").item()  # _ZERO = "0"
+
+    @property
+    @abstractmethod
+    def _prefix(self) -> "str": ...
+
+    def apply(
+        self,
+        df: pl.DataFrame,
+        to_col: Optional[str] = None,
+    ) -> pl.DataFrame:
+        if df.height == 0:
+            return df
+
+        query = pl.concat_str(
+            pl.lit(self._prefix),
+            pl.col(self._ID_COL).map_batches(
+                Base62Mapper._to_base62,
+                return_dtype=pl.String,
+                is_elementwise=True,
+            ),
+        )
+
+        if to_col is None:
+            to_col = self._ID_COL
+
+        return df.with_columns(query.alias(to_col))
+
+    @classmethod
+    def _to_base62(cls, int_col: pl.Series) -> pl.Series:
+        """Returns a series of dtype str with a base 62 representation of the integers in int_col.
+        The letters 0-9a-zA-Z are used as symbols for the representation.
+
+        Examples:
+
+            >>> import polars as pl
+            >>> s = pl.Series([0,10,20,60,53,66], dtype=pl.UInt32)
+            >>> Base62Mapper._to_base62(s).to_list()
+            ['0', 'a', 'k', 'Y', 'R', '14']
+
+            >>> s = pl.Series([0], dtype=pl.UInt32)
+            >>> Base62Mapper._to_base62(s).to_list()
+            ['0']
+        """
+        assert isinstance(
+            int_col.dtype, pl.UInt32
+        ), "_to_base62() only works for UInt32 id columns"
+
+        largest_id = int_col.max()
+        if largest_id == 0:
+            max_digits = 1
+        else:
+            max_digits = math.floor(math.log(largest_id, cls._BASE)) + 1  # type: ignore
+
+        digits = []
+
+        for i in range(max_digits):
+            remainder = int_col % cls._BASE
+
+            digits.append(
+                remainder.to_frame(name="code")
+                .join(cls._CHAR_TABLE, on="code", how="left")
+                .select("char")
+                .rename({"char": f"digit{i}"})
+            )
+            int_col //= cls._BASE
+
+        return (
+            pl.concat(reversed(digits), how="horizontal")
+            .select(pl.concat_str(pl.all()))
+            .to_series()
+            .str.strip_chars_start(cls._ZERO)
+            .replace("", cls._ZERO)
+        )
+
+    def _element_to_map(self, element) -> pl.DataFrame:
+        return self.apply(element.ids.select(self._ID_COL), to_col=NAME_COL)
+
+
+class Base62VarMapper(Base62Mapper):
+    """
+    Examples:
+        >>> import polars as pl
+        >>> from pyoframe import Model, Variable
+        >>> from pyoframe.constants import VAR_KEY
+        >>> m = Model()
+        >>> m.x = Variable(pl.DataFrame({"t": range(1,63)}))
+        >>> (m.x.filter(t=11)+1).to_str()
+        '[11]: 1  + x[11]'
+        >>> (m.x.filter(t=11)+1).to_str(var_map=Base62VarMapper(Variable))
+        '[11]: 1  + xb'
+
+        >>> Base62VarMapper(Variable).apply(pl.DataFrame({VAR_KEY: []}))
+        shape: (0, 1)
+        ┌───────────────┐
+        │ __variable_id │
+        │ ---           │
+        │ null          │
+        ╞═══════════════╡
+        └───────────────┘
+    """
+
+    @property
+    def _prefix(self) -> "str":
+        return "x"
+
+
+class Base62ConstMapper(Base62Mapper):
+
+    @property
+    def _prefix(self) -> "str":
+        return "c"

pyoframe/model.py

@@ -0,0 +1,92 @@
+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 pyoframe.objective import Objective
+from pyoframe.variables import Variable
+from pyoframe.io import to_file
+from pyoframe.solvers import solve
+
+
+class Model:
+    """
+    Represents a mathematical optimization model. Add variables, constraints, and an objective to the model by setting attributes.
+    """
+
+    def __init__(self, name=None):
+        self._variables: List[Variable] = []
+        self._constraints: List[Constraint] = []
+        self._objective: Optional[Objective] = None
+        self.var_map = (
+            NamedMapper(Variable) if Config.print_uses_variable_names else None
+        )
+        self.io_mappers: Optional[IOMappers] = None
+        self.name = name
+
+    @property
+    def variables(self) -> List[Variable]:
+        return self._variables
+
+    @property
+    def binary_variables(self) -> Iterable[Variable]:
+        return (v for v in self.variables if v.vtype == VType.BINARY)
+
+    @property
+    def integer_variables(self) -> Iterable[Variable]:
+        return (v for v in self.variables if v.vtype == VType.INTEGER)
+
+    @property
+    def constraints(self):
+        return self._constraints
+
+    @property
+    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 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:
+        return f"""Model '{self.name}' ({len(self.variables)} vars, {len(self.constraints)} constrs, {1 if self.objective else "no"} obj)"""
+
+    to_file = to_file
+    solve = solve

pyoframe/model_element.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import 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
+
+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
+        assert len(data.columns) == len(
+            set(data.columns)
+        ), "Duplicate column names found."
+
+        cols = _get_dimensions(data)
+        if cols is None:
+            cols = []
+        cols += [col for col in RESERVED_COL_KEYS if col in data.columns]
+
+        # Reorder columns to keep things consistent
+        data = data.select(cols)
+
+        # Cast to proper dtype
+        if COEF_KEY in data.columns:
+            data = data.cast({COEF_KEY: pl.Float64})
+        if VAR_KEY in data.columns:
+            data = data.cast({VAR_KEY: pl.UInt32})
+
+        self._data = data
+        self._model: Optional[Model] = None
+        self.name = None
+        super().__init__(**kwargs)
+
+    @property
+    def data(self) -> pl.DataFrame:
+        return self._data
+
+    @property
+    def dimensions(self) -> Optional[List[str]]:
+        """
+        The names of the data's dimensions.
+
+        Examples:
+            >>> from pyoframe.variables import Variable
+            >>> # A variable with no dimensions
+            >>> Variable().dimensions
+
+            >>> # A variable with dimensions of "hour" and "city"
+            >>> Variable([{"hour": ["00:00", "06:00", "12:00", "18:00"]}, {"city": ["Toronto", "Berlin", "Paris"]}]).dimensions
+            ['hour', 'city']
+        """
+        return _get_dimensions(self.data)
+
+    @property
+    def dimensions_unsafe(self) -> List[str]:
+        """
+        Same as `dimensions` but returns an empty list if there are no dimensions instead of None.
+        When unsure, use `dimensions` instead since the type checker forces users to handle the None case (no dimensions).
+        """
+        dims = self.dimensions
+        if dims is None:
+            return []
+        return dims
+
+    @property
+    def shape(self) -> Dict[str, int]:
+        """
+        The number of indices in each dimension.
+
+        Examples:
+            >>> from pyoframe.variables import Variable
+            >>> # A variable with no dimensions
+            >>> Variable().shape
+            {}
+            >>> # A variable with dimensions of "hour" and "city"
+            >>> Variable([{"hour": ["00:00", "06:00", "12:00", "18:00"]}, {"city": ["Toronto", "Berlin", "Paris"]}]).shape
+            {'hour': 4, 'city': 3}
+        """
+        dims = self.dimensions
+        if dims is None:
+            return {}
+        return {dim: self.data[dim].n_unique() for dim in dims}
+
+    def __len__(self) -> int:
+        dims = self.dimensions
+        if dims is None:
+            return 1
+        return self.data.select(dims).n_unique()
+
+    @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")

pyoframe/monkey_patch.py

@@ -0,0 +1,54 @@
+import pandas as pd
+import polars as pl
+from pyoframe.constraints import SupportsMath
+from pyoframe.constraints import Expression
+from functools import wraps
+
+from pyoframe.constants import COEF_KEY, CONST_TERM, VAR_KEY
+
+# pyright: reportAttributeAccessIssue=false
+
+
+def _patch_class(cls):
+    def _patch_method(func):
+        @wraps(func)
+        def wrapper(self, other):
+            if isinstance(other, SupportsMath):
+                return NotImplemented
+            return func(self, other)
+
+        return wrapper
+
+    cls.__add__ = _patch_method(cls.__add__)
+    cls.__mul__ = _patch_method(cls.__mul__)
+    cls.__sub__ = _patch_method(cls.__sub__)
+    cls.__le__ = _patch_method(cls.__le__)
+    cls.__ge__ = _patch_method(cls.__ge__)
+    cls.__contains__ = _patch_method(cls.__contains__)
+
+
+def _dataframe_to_expr(self: pl.DataFrame) -> Expression:
+    return Expression(
+        self.rename({self.columns[-1]: COEF_KEY})
+        .drop_nulls(COEF_KEY)
+        .with_columns(pl.lit(CONST_TERM).alias(VAR_KEY))
+    )
+
+
+def patch_dataframe_libraries():
+    """
+    Applies two patches to the DataFrame and Series classes of both pandas and polars.
+    1) Patches arithmetic operators (e.g. `__add__`) such that operations between DataFrames/Series and `Expressionable`s
+        are not supported (i.e. `return NotImplemented`). This leads Python to try the reverse operation (e.g. `__radd__`)
+        which is supported by the `Expressionable` class.
+    2) Adds a `to_expr` method to DataFrame/Series that allows them to be converted to an `Expression` object.
+        Series become dataframes and dataframes become expressions where everything but the last column are treated as dimensions.
+    """
+    _patch_class(pd.DataFrame)
+    _patch_class(pd.Series)
+    _patch_class(pl.DataFrame)
+    _patch_class(pl.Series)
+    pl.DataFrame.to_expr = _dataframe_to_expr
+    pl.Series.to_expr = lambda self: self.to_frame().to_expr()
+    pd.DataFrame.to_expr = lambda self: pl.from_pandas(self).to_expr()
+    pd.Series.to_expr = lambda self: self.to_frame().reset_index().to_expr()

pyoframe/objective.py

@@ -0,0 +1,42 @@
+from typing import Optional, Union
+from pyoframe.constants import ObjSense, ObjSenseValue
+from pyoframe.constraints import SupportsMath, Expression
+
+
+class Objective(Expression):
+    r"""
+    Examples:
+        >>> from pyoframe import Variable, Model, sum
+        >>> m = Model()
+        >>> m.a = Variable()
+        >>> m.b = Variable({"dim1": [1, 2, 3]})
+        >>> m.maximize = m.a + sum("dim1", m.b)
+        >>> m.maximize
+        <Objective size=1 dimensions={} terms=4>
+        maximize: a + b[1] + b[2] + b[3]
+    """
+
+    def __init__(
+        self, expr: SupportsMath, sense: Union[ObjSense, ObjSenseValue]
+    ) -> None:
+        self.sense = ObjSense(sense)
+
+        expr = expr.to_expr()
+        super().__init__(expr.to_expr().data)
+        self._model = expr._model
+        assert (
+            self.dimensions is None
+        ), "Objective cannot have dimensions as it must be a single expression"
+        self._value: Optional[float] = None
+
+    @property
+    def value(self):
+        if self._value is None:
+            raise ValueError(
+                "Objective value is not available before solving the model"
+            )
+        return self._value
+
+    @value.setter
+    def value(self, value):
+        self._value = value