pyoframe 0.0.4__py3-none-any.whl

pyoframe/__init__.py

@@ -0,0 +1,15 @@
+"""
+Pyoframe's public API.
+Also applies the monkey patch to the DataFrame libraries.
+"""
+
+from pyoframe.monkey_patch import patch_dataframe_libraries
+from pyoframe.constraints import sum, sum_by, Set, Constraint
+from pyoframe.constants import Config
+from pyoframe.variables import Variable
+from pyoframe.model import Model
+from pyoframe.constants import VType
+
+patch_dataframe_libraries()
+
+__all__ = ["sum", "sum_by", "Variable", "Model", "Set", "VType", "Config", "Constraint"]

pyoframe/_arithmetic.py

@@ -0,0 +1,228 @@
+from typing import TYPE_CHECKING, List, Optional
+import polars as pl
+
+from pyoframe.constants import (
+    COEF_KEY,
+    RESERVED_COL_KEYS,
+    VAR_KEY,
+    UnmatchedStrategy,
+    Config,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from pyoframe.constraints import Expression
+
+
+class PyoframeError(Exception):
+    pass
+
+
+def _add_expressions(*expressions: "Expression") -> "Expression":
+    try:
+        return _add_expressions_core(*expressions)
+    except PyoframeError as error:
+        raise PyoframeError(
+            "Failed to add expressions:\n"
+            + " + ".join(
+                e.to_str(include_header=True, include_data=False) for e in expressions
+            )
+            + "\nDue to error:\n"
+            + str(error)
+        ) from error
+
+
+def _add_expressions_core(*expressions: "Expression") -> "Expression":
+    # Mapping of how a sum of two expressions should propogate the unmatched strategy
+    propogatation_strategies = {
+        (UnmatchedStrategy.DROP, UnmatchedStrategy.DROP): UnmatchedStrategy.DROP,
+        (
+            UnmatchedStrategy.UNSET,
+            UnmatchedStrategy.UNSET,
+        ): UnmatchedStrategy.UNSET,
+        (UnmatchedStrategy.KEEP, UnmatchedStrategy.KEEP): UnmatchedStrategy.KEEP,
+        (UnmatchedStrategy.DROP, UnmatchedStrategy.KEEP): UnmatchedStrategy.UNSET,
+        (UnmatchedStrategy.DROP, UnmatchedStrategy.UNSET): UnmatchedStrategy.DROP,
+        (UnmatchedStrategy.KEEP, UnmatchedStrategy.UNSET): UnmatchedStrategy.KEEP,
+    }
+
+    assert len(expressions) > 1, "Need at least two expressions to add together."
+
+    dims = expressions[0].dimensions
+
+    if dims is None:
+        requires_join = False
+        dims = []
+    elif Config.disable_unmatched_checks:
+        requires_join = any(
+            expr.unmatched_strategy
+            not in (UnmatchedStrategy.KEEP, UnmatchedStrategy.UNSET)
+            for expr in expressions
+        )
+    else:
+        requires_join = any(
+            expr.unmatched_strategy != UnmatchedStrategy.KEEP for expr in expressions
+        )
+
+    has_dim_conflict = any(
+        sorted(dims) != sorted(expr.dimensions_unsafe) for expr in expressions[1:]
+    )
+
+    # If we cannot use .concat compute the sum in a pairwise manner
+    if len(expressions) > 2 and (has_dim_conflict or requires_join):
+        result = expressions[0]
+        for expr in expressions[1:]:
+            result = _add_expressions_core(result, expr)
+        return result
+
+    if has_dim_conflict:
+        assert len(expressions) == 2
+        expressions = (
+            _add_dimension(expressions[0], expressions[1]),
+            _add_dimension(expressions[1], expressions[0]),
+        )
+        assert sorted(expressions[0].dimensions_unsafe) == sorted(
+            expressions[1].dimensions_unsafe
+        )
+
+    dims = expressions[0].dimensions_unsafe
+    # Check no dims conflict
+    assert all(
+        sorted(dims) == sorted(expr.dimensions_unsafe) for expr in expressions[1:]
+    )
+    if requires_join:
+        assert len(expressions) == 2
+        assert dims != []
+        left, right = expressions[0], expressions[1]
+
+        # Order so that drop always comes before keep, and keep always comes before default
+        if (left.unmatched_strategy, right.unmatched_strategy) in (
+            (UnmatchedStrategy.UNSET, UnmatchedStrategy.DROP),
+            (UnmatchedStrategy.UNSET, UnmatchedStrategy.KEEP),
+            (UnmatchedStrategy.KEEP, UnmatchedStrategy.DROP),
+        ):
+            left, right = right, left
+
+        def get_indices(expr):
+            return expr.data.select(dims).unique(maintain_order=True)
+
+        left_data, right_data = left.data, right.data
+
+        strat = (left.unmatched_strategy, right.unmatched_strategy)
+
+        propogate_strat = propogatation_strategies[strat]
+
+        if strat == (UnmatchedStrategy.DROP, UnmatchedStrategy.DROP):
+            left_data = left.data.join(get_indices(right), how="inner", on=dims)
+            right_data = right.data.join(get_indices(left), how="inner", on=dims)
+        elif strat == (UnmatchedStrategy.UNSET, UnmatchedStrategy.UNSET):
+            assert (
+                not Config.disable_unmatched_checks
+            ), "This code should not be reached when unmatched checks are disabled."
+            outer_join = get_indices(left).join(
+                get_indices(right), how="outer", on=dims
+            )
+            if outer_join.get_column(dims[0]).null_count() > 0:
+                raise PyoframeError(
+                    "Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()\n"
+                    + str(outer_join.filter(outer_join.get_column(dims[0]).is_null()))
+                )
+            if outer_join.get_column(dims[0] + "_right").null_count() > 0:
+                raise PyoframeError(
+                    "Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()\n"
+                    + str(
+                        outer_join.filter(
+                            outer_join.get_column(dims[0] + "_right").is_null()
+                        )
+                    )
+                )
+        elif strat == (UnmatchedStrategy.DROP, UnmatchedStrategy.KEEP):
+            left_data = get_indices(right).join(left.data, how="left", on=dims)
+        elif strat == (UnmatchedStrategy.DROP, UnmatchedStrategy.UNSET):
+            left_data = get_indices(right).join(left.data, how="left", on=dims)
+            if left_data.get_column(COEF_KEY).null_count() > 0:
+                raise PyoframeError(
+                    "Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()\n"
+                    + str(left_data.filter(left_data.get_column(COEF_KEY).is_null()))
+                )
+        elif strat == (UnmatchedStrategy.KEEP, UnmatchedStrategy.UNSET):
+            assert (
+                not Config.disable_unmatched_checks
+            ), "This code should not be reached when unmatched checks are disabled."
+            unmatched = right.data.join(get_indices(left), how="anti", on=dims)
+            if len(unmatched) > 0:
+                raise PyoframeError(
+                    "Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()\n"
+                    + str(unmatched)
+                )
+        else:  # pragma: no cover
+            assert False, "This code should've never been reached!"
+
+        expr_data = [left_data, right_data]
+    else:
+        propogate_strat = expressions[0].unmatched_strategy
+        expr_data = [expr.data for expr in expressions]
+
+    # Sort columns to allow for concat
+    expr_data = [e.select(sorted(e.columns)) for e in expr_data]
+
+    data = pl.concat(expr_data, how="vertical_relaxed")
+    data = data.group_by(dims + [VAR_KEY], maintain_order=True).sum()
+
+    new_expr = expressions[0]._new(data)
+    new_expr.unmatched_strategy = propogate_strat
+
+    return new_expr
+
+
+def _add_dimension(self: "Expression", target: "Expression") -> "Expression":
+    target_dims = target.dimensions
+    if target_dims is None:
+        return self
+    dims = self.dimensions
+    if dims is None:
+        dims_in_common = []
+        missing_dims = target_dims
+    else:
+        dims_in_common = [dim for dim in dims if dim in target_dims]
+        missing_dims = [dim for dim in target_dims if dim not in dims]
+
+    # We're already at the size of our target
+    if not missing_dims:
+        return self
+
+    if not set(missing_dims) <= set(self.allowed_new_dims):
+        raise PyoframeError(
+            f"Dataframe has missing dimensions {missing_dims}. If this is intentional, use .add_dim()\n{self.data}"
+        )
+
+    target_data = target.data.select(target_dims).unique(maintain_order=True)
+
+    if not dims_in_common:
+        return self._new(self.data.join(target_data, how="cross"))
+
+    # If drop, we just do an inner join to get into the shape of the other
+    if self.unmatched_strategy == UnmatchedStrategy.DROP:
+        return self._new(self.data.join(target_data, on=dims_in_common, how="inner"))
+
+    result = self.data.join(target_data, on=dims_in_common, how="left")
+    right_has_missing = result.get_column(missing_dims[0]).null_count() > 0
+    if right_has_missing:
+        raise PyoframeError(
+            f"Cannot add dimension {missing_dims} since it contains unmatched values. If this is intentional, consider using .drop_unmatched()"
+        )
+    return self._new(result)
+
+
+def _get_dimensions(df: pl.DataFrame) -> Optional[List[str]]:
+    """
+    Returns the dimensions of the DataFrame. Reserved columns do not count as dimensions.
+    If there are no dimensions, returns None to force caller to handle this special case.
+
+    Examples:
+        >>> import polars as pl
+        >>> _get_dimensions(pl.DataFrame({"x": [1, 2, 3], "y": [1, 2, 3]}))
+        ['x', 'y']
+        >>> _get_dimensions(pl.DataFrame({"__variable_id": [1, 2, 3]}))
+    """
+    result = [col for col in df.columns if col not in RESERVED_COL_KEYS]
+    return result if result else None

pyoframe/constants.py

@@ -0,0 +1,280 @@
+"""
+File containing shared constants used across the package.
+
+Code is heavily based on the `linopy` package by Fabian Hofmann.
+
+MIT License
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+import typing
+from typing import Any, Literal, Optional, Union
+import polars as pl
+
+
+COEF_KEY = "__coeff"
+VAR_KEY = "__variable_id"
+CONSTRAINT_KEY = "__constraint_id"
+SOLUTION_KEY = "solution"
+DUAL_KEY = "dual"
+NAME_COL = "__name"
+
+CONST_TERM = 0
+
+RESERVED_COL_KEYS = (
+    COEF_KEY,
+    VAR_KEY,
+    CONSTRAINT_KEY,
+    SOLUTION_KEY,
+    DUAL_KEY,
+    NAME_COL,
+)
+
+
+class _ConfigMeta(type):
+    """Metaclass for Config that stores the default values of all configuration options."""
+
+    def __init__(cls, name, bases, dct):
+        super().__init__(name, bases, dct)
+        cls._defaults = {
+            k: v
+            for k, v in dct.items()
+            if not k.startswith("_") and type(v) != classmethod
+        }
+
+
+class Config(metaclass=_ConfigMeta):
+    disable_unmatched_checks: bool = False
+    print_float_precision: Optional[int] = 5
+    print_uses_variable_names: bool = True
+
+    @classmethod
+    def reset_defaults(cls):
+        """
+        Resets all configuration options to their default values.
+        """
+        for key, value in cls._defaults.items():
+            setattr(cls, key, value)
+
+
+class ConstraintSense(Enum):
+    LE = "<="
+    GE = ">="
+    EQ = "="
+
+
+class ObjSense(Enum):
+    MIN = "minimize"
+    MAX = "maximize"
+
+
+class VType(Enum):
+    CONTINUOUS = "continuous"
+    BINARY = "binary"
+    INTEGER = "integer"
+
+
+class UnmatchedStrategy(Enum):
+    UNSET = "not_set"
+    DROP = "drop"
+    KEEP = "keep"
+
+
+# This is a hack to get the Literal type for VType
+# See: https://stackoverflow.com/questions/67292470/type-hinting-enum-member-value-in-python
+ObjSenseValue = Literal["minimize", "maximize"]
+VTypeValue = Literal["continuous", "binary", "integer"]
+for enum, type in [(ObjSense, ObjSenseValue), (VType, VTypeValue)]:
+    assert set(typing.get_args(type)) == {vtype.value for vtype in enum}
+
+
+class ModelStatus(Enum):
+    """
+    Model status.
+
+    The set of possible model status is a superset of the solver status
+    set.
+    """
+
+    ok = "ok"
+    warning = "warning"
+    error = "error"
+    aborted = "aborted"
+    unknown = "unknown"
+    initialized = "initialized"
+
+
+class SolverStatus(Enum):
+    """
+    Solver status.
+    """
+
+    ok = "ok"
+    warning = "warning"
+    error = "error"
+    aborted = "aborted"
+    unknown = "unknown"
+
+    @classmethod
+    def process(cls, status: str) -> "SolverStatus":
+        try:
+            return cls(status)
+        except ValueError:
+            return cls("unknown")
+
+    @classmethod
+    def from_termination_condition(
+        cls, termination_condition: "TerminationCondition"
+    ) -> "SolverStatus":
+        for status in STATUS_TO_TERMINATION_CONDITION_MAP:
+            if termination_condition in STATUS_TO_TERMINATION_CONDITION_MAP[status]:
+                return status
+        return cls("unknown")
+
+
+class TerminationCondition(Enum):
+    """
+    Termination condition of the solver.
+    """
+
+    # UNKNOWN
+    unknown = "unknown"
+
+    # OK
+    optimal = "optimal"
+    time_limit = "time_limit"
+    iteration_limit = "iteration_limit"
+    terminated_by_limit = "terminated_by_limit"
+    suboptimal = "suboptimal"
+
+    # WARNING
+    unbounded = "unbounded"
+    infeasible = "infeasible"
+    infeasible_or_unbounded = "infeasible_or_unbounded"
+    other = "other"
+
+    # ERROR
+    internal_solver_error = "internal_solver_error"
+    error = "error"
+
+    # ABORTED
+    user_interrupt = "user_interrupt"
+    resource_interrupt = "resource_interrupt"
+    licensing_problems = "licensing_problems"
+
+    @classmethod
+    def process(
+        cls, termination_condition: Union[str, "TerminationCondition"]
+    ) -> "TerminationCondition":
+        try:
+            return cls(termination_condition)
+        except ValueError:
+            return cls("unknown")
+
+
+STATUS_TO_TERMINATION_CONDITION_MAP = {
+    SolverStatus.ok: [
+        TerminationCondition.optimal,
+        TerminationCondition.iteration_limit,
+        TerminationCondition.time_limit,
+        TerminationCondition.terminated_by_limit,
+        TerminationCondition.suboptimal,
+    ],
+    SolverStatus.warning: [
+        TerminationCondition.unbounded,
+        TerminationCondition.infeasible,
+        TerminationCondition.infeasible_or_unbounded,
+        TerminationCondition.other,
+    ],
+    SolverStatus.error: [
+        TerminationCondition.internal_solver_error,
+        TerminationCondition.error,
+    ],
+    SolverStatus.aborted: [
+        TerminationCondition.user_interrupt,
+        TerminationCondition.resource_interrupt,
+        TerminationCondition.licensing_problems,
+    ],
+    SolverStatus.unknown: [TerminationCondition.unknown],
+}
+
+
+@dataclass
+class Status:
+    """
+    Status and termination condition of the solver.
+    """
+
+    status: SolverStatus
+    termination_condition: TerminationCondition
+
+    @classmethod
+    def process(cls, status: str, termination_condition: str) -> "Status":
+        return cls(
+            status=SolverStatus.process(status),
+            termination_condition=TerminationCondition.process(termination_condition),
+        )
+
+    @classmethod
+    def from_termination_condition(
+        cls, termination_condition: Union["TerminationCondition", str]
+    ) -> "Status":
+        termination_condition = TerminationCondition.process(termination_condition)
+        solver_status = SolverStatus.from_termination_condition(termination_condition)
+        return cls(solver_status, termination_condition)
+
+    @property
+    def is_ok(self) -> bool:
+        return self.status == SolverStatus.ok
+
+
+@dataclass
+class Solution:
+    """
+    Solution returned by the solver.
+    """
+
+    primal: pl.DataFrame
+    dual: Optional[pl.DataFrame]
+    objective: float
+
+
+@dataclass
+class Result:
+    """
+    Result of the optimization.
+    """
+
+    status: Status
+    solution: Optional[Solution] = None
+    solver_model: Optional[Any] = None
+
+    def __repr__(self) -> str:
+        solver_model_string = (
+            "not available" if self.solver_model is None else "available"
+        )
+
+        res = (
+            f"Status: {self.status.status.value}\n"
+            f"Termination condition: {self.status.termination_condition.value}\n"
+        )
+        if self.solution is not None:
+            res += (
+                f"Solution: {len(self.solution.primal)} primals, {len(self.solution.dual) if self.solution.dual is not None else 0} duals\n"
+                f"Objective: {self.solution.objective:.2e}\n"
+            )
+        res += f"Solver model: {solver_model_string}\n"
+
+        return res
+
+    def info(self):
+        status = self.status
+
+        if status.is_ok:
+            if status.termination_condition == TerminationCondition.suboptimal:
+                print(f"Optimization solution is sub-optimal: \n{self}\n")
+            else:
+                print(f" Optimization successful: \n{self}\n")
+        else:
+            print(f"Optimization failed: \n{self}\n")