pyoframe 0.2.0__py3-none-any.whl → 1.0.0a0__py3-none-any.whl

pyoframe/_model_element.py

@@ -0,0 +1,175 @@
+"""Defines the base classes used in Pyoframe."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+from pyoframe._arithmetic import _get_dimensions
+from pyoframe._constants import (
+    COEF_KEY,
+    KEY_TYPE,
+    QUAD_VAR_KEY,
+    RESERVED_COL_KEYS,
+    VAR_KEY,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pyoframe import Model
+
+
+class ModelElement(ABC):
+    """The base class for elements of a Model such as [][pyoframe.Variable] and [][pyoframe.Constraint]."""
+
+    def __init__(self, data: pl.DataFrame, name="unnamed") -> 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: KEY_TYPE})
+        if QUAD_VAR_KEY in data.columns:
+            data = data.cast({QUAD_VAR_KEY: KEY_TYPE})
+
+        self._data = data
+        self._model: Model | None = None
+        self.name: str = name  # gets overwritten if object is added to model
+
+    def _on_add_to_model(self, model: Model, name: str):
+        self.name = name
+        self._model = model
+
+    @property
+    def data(self) -> pl.DataFrame:
+        """Returns the object's underlying Polars DataFrame."""
+        return self._data
+
+    @property
+    def dimensions(self) -> list[str] | None:
+        """The names of the data's dimensions.
+
+        Examples:
+            A variable with no dimensions
+            >>> pf.Variable().dimensions
+
+            A variable with dimensions of "hour" and "city"
+            >>> pf.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:
+            A variable with no dimensions
+            >>> pf.Variable().shape
+            {}
+
+            A variable with dimensions of "hour" and "city"
+            >>> pf.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 estimated_size(self, unit: pl.SizeUnit = "b") -> int | float:
+        """Returns the estimated size of the object in bytes.
+
+        Only considers the size of the underlying DataFrame(s) since other components (e.g., the object name) are negligible.
+
+        Parameters:
+            unit:
+                See [`polars.DataFrame.estimated_size`](https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.estimated_size.html).
+
+        Examples:
+            >>> m = pf.Model()
+
+            A dimensionless variable contains just a 32 bit (4 bytes) unsigned integer (the variable ID).
+
+            >>> m.x = pf.Variable()
+            >>> m.x.estimated_size()
+            4
+
+            A dimensioned variable contains, for every row, a 32 bit ID and, in this case, a 64 bit `dim_x` value (1200 bytes total).
+
+            >>> m.y = pf.Variable(pf.Set(dim_x=range(100)))
+            >>> m.y.estimated_size()
+            1200
+        """
+        return self.data.estimated_size(unit)
+
+    def _add_shape_to_columns(self, df: pl.DataFrame) -> pl.DataFrame:
+        """Adds the shape of the data to the columns of the DataFrame.
+
+        This is used for displaying the shape in the string representation of the object.
+        """
+        shape = self.shape
+        return df.rename(lambda col: f"{col}\n({shape[col]})" if col in shape else col)
+
+    def __len__(self) -> int:
+        dims = self.dimensions
+        if dims is None:
+            return 1
+        return self.data.select(dims).n_unique()
+
+
+class ModelElementWithId(ModelElement):
+    """Extends ModelElement with 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.
+    """
+
+    @property
+    def _has_ids(self) -> bool:
+        return self._get_id_column_name() in self.data.columns
+
+    def _assert_has_ids(self):
+        if not self._has_ids:
+            raise ValueError(
+                f"Cannot use '{self.__class__.__name__}' before it has been added to a model."
+            )
+
+    @classmethod
+    @abstractmethod
+    def _get_id_column_name(cls) -> str:
+        """Returns the name of the column containing the IDs."""

pyoframe/_monkey_patch.py

@@ -0,0 +1,80 @@
+"""Defines the functions used to monkey patch polars and pandas."""
+
+from functools import wraps
+
+import pandas as pd
+import polars as pl
+
+from pyoframe._constants import COEF_KEY, CONST_TERM, VAR_KEY
+from pyoframe._core import Expression, SupportsMath
+
+
+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 polars_df_to_expr(self: pl.DataFrame) -> Expression:
+    """Converts a [polars](https://pola.rs/) `DataFrame` to a Pyoframe [Expression][pyoframe.Expression] by using the last column for values and the previous columns as dimensions.
+
+    See [Special Functions](../learn/concepts/special-functions.md#dataframeto_expr) for more details.
+
+    Examples:
+        >>> import polars as pl
+        >>> df = pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
+        >>> df.to_expr()
+        <Expression height=3 terms=3 type=constant>
+        ┌─────┬─────┬────────────┐
+        │ x   ┆ y   ┆ expression │
+        │ (3) ┆ (3) ┆            │
+        ╞═════╪═════╪════════════╡
+        │ 1   ┆ 4   ┆ 7          │
+        │ 2   ┆ 5   ┆ 8          │
+        │ 3   ┆ 6   ┆ 9          │
+        └─────┴─────┴────────────┘
+    """
+    name = self.columns[-1]
+    return Expression(
+        self.rename({name: COEF_KEY})
+        .drop_nulls(COEF_KEY)
+        .with_columns(pl.lit(CONST_TERM).alias(VAR_KEY)),
+        name=name,
+    )
+
+
+def pandas_df_to_expr(self: pd.DataFrame) -> Expression:
+    """Same as [`polars.DataFrame.to_expr`](./polars.DataFrame.to_expr.md), but for [pandas](https://pandas.pydata.org/) DataFrames."""
+    return polars_df_to_expr(pl.from_pandas(self))
+
+
+def patch_dataframe_libraries():
+    """Patches 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 = polars_df_to_expr
+    pd.DataFrame.to_expr = pandas_df_to_expr
+    # TODO make a set instead!
+    pl.Series.to_expr = lambda self: self.to_frame().to_expr()
+    pd.Series.to_expr = lambda self: self.to_frame().reset_index().to_expr()

pyoframe/{objective.py → _objective.py}

@@ -1,12 +1,17 @@
+"""Defines the `Objective` class for optimization models."""
+
 from __future__ import annotations
 
 import pyoptinterface as poi
 
-from pyoframe.core import Expression, SupportsToExpr
+from pyoframe._constants import ObjSense
+from pyoframe._core import Expression, SupportsToExpr
 
 
+# TODO don't subclass Expression to avoid a bunch of unnecessary functions being available.
 class Objective(Expression):
-    """
+    """The objective for an optimization model.
+
     Examples:
         An `Objective` is automatically created when an `Expression` is assigned to `.minimize` or `.maximize`
 
@@ -15,8 +20,8 @@ class Objective(Expression):
         >>> m.con = m.A + m.B <= 10
         >>> m.maximize = 2 * m.B + 4
         >>> m.maximize
-        <Objective size=1 dimensions={} terms=2>
-        objective: 2 B +4
+        <Objective terms=2 type=linear>
+        2 B +4
 
         The objective value can be retrieved with from the solver once the model is solved using `.value`.
 
@@ -38,11 +43,13 @@ class Objective(Expression):
         Objectives cannot be created from dimensioned expressions since an objective must be a single expression.
 
         >>> m = pf.Model()
-        >>> m.dimensioned_variable = pf.Variable({"city": ["Toronto", "Berlin", "Paris"]})
+        >>> m.dimensioned_variable = pf.Variable(
+        ...     {"city": ["Toronto", "Berlin", "Paris"]}
+        ... )
         >>> m.maximize = m.dimensioned_variable
         Traceback (most recent call last):
         ...
-        ValueError: Objective cannot be created from a dimensioned expression. Did you forget to use pf.sum()?
+        ValueError: Objective cannot be created from a dimensioned expression. Did you forget to use .sum()?
 
         Objectives cannot be overwritten.
 
@@ -63,30 +70,58 @@ class Objective(Expression):
             expr = Expression.constant(expr)
         else:
             expr = expr.to_expr()
-        super().__init__(expr.data)
+        super().__init__(expr.data, name="objective")
         self._model = expr._model
         if self.dimensions is not None:
             raise ValueError(
-                "Objective cannot be created from a dimensioned expression. Did you forget to use pf.sum()?"
+                "Objective cannot be created from a dimensioned expression. Did you forget to use .sum()?"
             )
 
     @property
     def value(self) -> float:
-        """
-        The value of the objective function (only available after solving the model).
+        """The value of the objective function (only available after solving the model).
 
         This value is obtained by directly querying the solver.
         """
-        return self._model.poi.get_model_attribute(poi.ModelAttribute.ObjectiveValue)
+        assert self._model is not None, (
+            "Objective must be part of a model before it is queried."
+        )
+
+        if (
+            self._model.attr.TerminationStatus
+            == poi.TerminationStatusCode.OPTIMIZE_NOT_CALLED
+        ):
+            raise ValueError(
+                "Cannot retrieve the objective value before calling model.optimize()."
+            )
+
+        obj_value: float = self._model.attr.ObjectiveValue
+        if (
+            not self._model.solver.supports_objective_sense
+            and self._model.sense == ObjSense.MAX
+        ):
+            obj_value *= -1
+        return obj_value
 
-    def on_add_to_model(self, model, name):
-        super().on_add_to_model(model, name)
+    def _on_add_to_model(self, model, name):
+        super()._on_add_to_model(model, name)
         assert self._model is not None
         if self._model.sense is None:
             raise ValueError(
                 "Can't set an objective without specifying the sense. Did you use .objective instead of .minimize or .maximize ?"
             )
-        self._model.poi.set_objective(self.to_poi(), sense=self._model.sense.to_poi())
+
+        kwargs = {}
+        if (
+            not self._model.solver.supports_objective_sense
+            and self._model.sense == ObjSense.MAX
+        ):
+            poi_expr = (-self)._to_poi()
+            kwargs["sense"] = poi.ObjectiveSense.Minimize
+        else:
+            poi_expr = self._to_poi()
+            kwargs["sense"] = self._model.sense._to_poi()
+        self._model.poi.set_objective(poi_expr, **kwargs)
 
     def __iadd__(self, other):
         return Objective(self + other, _constructive=True)