pyoframe 0.0.11__py3-none-any.whl → 0.1.1__py3-none-any.whl

pyoframe/__init__.py

@@ -3,11 +3,10 @@ Pyoframe's public API.
 Also applies the monkey patch to the DataFrame libraries.
 """
 
-from pyoframe.monkey_patch import patch_dataframe_libraries
-from pyoframe.core import sum, sum_by, Set, Constraint, Expression, Variable
-from pyoframe.constants import Config
+from pyoframe.constants import Config, VType
+from pyoframe.core import Constraint, Expression, Set, Variable, sum, sum_by
 from pyoframe.model import Model
-from pyoframe.constants import VType
+from pyoframe.monkey_patch import patch_dataframe_libraries
 
 patch_dataframe_libraries()
 

pyoframe/_arithmetic.py

@@ -1,19 +1,63 @@
+"""
+Defines helper functions for doing arithmetic operations on expressions (e.g. addition).
+"""
+
 from typing import TYPE_CHECKING, List, Optional
+
 import polars as pl
 
 from pyoframe.constants import (
     COEF_KEY,
+    CONST_TERM,
+    KEY_TYPE,
+    POLARS_VERSION,
+    QUAD_VAR_KEY,
     RESERVED_COL_KEYS,
     VAR_KEY,
-    UnmatchedStrategy,
     Config,
     PyoframeError,
+    UnmatchedStrategy,
 )
 
 if TYPE_CHECKING:  # pragma: no cover
     from pyoframe.core import Expression
 
 
+def _multiply_expressions(self: "Expression", other: "Expression") -> "Expression":
+    """
+    Multiply two or more expressions together.
+
+    Examples:
+        >>> import pyoframe as pf
+        >>> m = pf.Model("min")
+        >>> m.x1 = pf.Variable()
+        >>> m.x2 = pf.Variable()
+        >>> m.x3 = pf.Variable()
+        >>> result = 5 * m.x1 * m.x2
+        >>> result
+        <Expression size=1 dimensions={} terms=1 degree=2>
+        5 x2 * x1
+        >>> result * m.x3
+        Traceback (most recent call last):
+        ...
+        pyoframe.constants.PyoframeError: Failed to multiply expressions:
+        <Expression size=1 dimensions={} terms=1 degree=2> * <Expression size=1 dimensions={} terms=1>
+        Due to error:
+        Cannot multiply a quadratic expression by a non-constant.
+    """
+    try:
+        return _multiply_expressions_core(self, other)
+    except PyoframeError as error:
+        raise PyoframeError(
+            "Failed to multiply expressions:\n"
+            + " * ".join(
+                e.to_str(include_header=True, include_data=False) for e in [self, other]
+            )
+            + "\nDue to error:\n"
+            + str(error)
+        ) from error
+
+
 def _add_expressions(*expressions: "Expression") -> "Expression":
     try:
         return _add_expressions_core(*expressions)
@@ -28,6 +72,98 @@ def _add_expressions(*expressions: "Expression") -> "Expression":
         ) from error
 
 
+def _multiply_expressions_core(self: "Expression", other: "Expression") -> "Expression":
+    self_degree, other_degree = self.degree(), other.degree()
+    if self_degree + other_degree > 2:
+        # We know one of the two must be a quadratic since 1 + 1 is not greater than 2.
+        raise PyoframeError("Cannot multiply a quadratic expression by a non-constant.")
+    if self_degree < other_degree:
+        self, other = other, self
+        self_degree, other_degree = other_degree, self_degree
+    if other_degree == 1:
+        assert (
+            self_degree == 1
+        ), "This should always be true since the sum of degrees must be <=2."
+        return _quadratic_multiplication(self, other)
+
+    assert (
+        other_degree == 0
+    ), "This should always be true since other cases have already been handled."
+    multiplier = other.data.drop(
+        VAR_KEY
+    )  # QUAD_VAR_KEY doesn't need to be dropped since we know it doesn't exist
+
+    dims = self.dimensions_unsafe
+    other_dims = other.dimensions_unsafe
+    dims_in_common = [dim for dim in dims if dim in other_dims]
+
+    data = (
+        self.data.join(
+            multiplier,
+            on=dims_in_common if len(dims_in_common) > 0 else None,
+            how="inner" if dims_in_common else "cross",
+        )
+        .with_columns(pl.col(COEF_KEY) * pl.col(COEF_KEY + "_right"))
+        .drop(COEF_KEY + "_right")
+    )
+
+    return self._new(data)
+
+
+def _quadratic_multiplication(self: "Expression", other: "Expression") -> "Expression":
+    """
+    Multiply two expressions of degree 1.
+
+    Examples:
+        >>> import polars as pl
+        >>> df = pl.DataFrame({"dim": [1, 2, 3], "value": [1, 2, 3]})
+        >>> m = pf.Model()
+        >>> m.x1 = pf.Variable()
+        >>> m.x2 = pf.Variable()
+        >>> expr1 = df * m.x1
+        >>> expr2 = df * m.x2 * 2 + 4
+        >>> expr1 * expr2
+        <Expression size=3 dimensions={'dim': 3} terms=6 degree=2>
+        [1]: 4 x1 +2 x2 * x1
+        [2]: 8 x1 +8 x2 * x1
+        [3]: 12 x1 +18 x2 * x1
+        >>> (expr1 * expr2) - df * m.x1 * df * m.x2 * 2
+        <Expression size=3 dimensions={'dim': 3} terms=3>
+        [1]: 4 x1
+        [2]: 8 x1
+        [3]: 12 x1
+    """
+    dims = self.dimensions_unsafe
+    other_dims = other.dimensions_unsafe
+    dims_in_common = [dim for dim in dims if dim in other_dims]
+
+    data = (
+        self.data.join(
+            other.data,
+            on=dims_in_common if len(dims_in_common) > 0 else None,
+            how="inner" if dims_in_common else "cross",
+        )
+        .with_columns(pl.col(COEF_KEY) * pl.col(COEF_KEY + "_right"))
+        .drop(COEF_KEY + "_right")
+        .rename({VAR_KEY + "_right": QUAD_VAR_KEY})
+        # Swap VAR_KEY and QUAD_VAR_KEY so that VAR_KEy is always the larger one
+        .with_columns(
+            pl.when(pl.col(VAR_KEY) < pl.col(QUAD_VAR_KEY))
+            .then(pl.col(QUAD_VAR_KEY))
+            .otherwise(pl.col(VAR_KEY))
+            .alias(VAR_KEY),
+            pl.when(pl.col(VAR_KEY) < pl.col(QUAD_VAR_KEY))
+            .then(pl.col(VAR_KEY))
+            .otherwise(pl.col(QUAD_VAR_KEY))
+            .alias(QUAD_VAR_KEY),
+        )
+    )
+
+    data = _sum_like_terms(data)
+
+    return self._new(data)
+
+
 def _add_expressions_core(*expressions: "Expression") -> "Expression":
     # Mapping of how a sum of two expressions should propogate the unmatched strategy
     propogatation_strategies = {
@@ -116,7 +252,9 @@ def _add_expressions_core(*expressions: "Expression") -> "Expression":
                 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
+                get_indices(right),
+                how="full" if POLARS_VERSION.major >= 1 else "outer",
+                on=dims,
             )
             if outer_join.get_column(dims[0]).null_count() > 0:
                 raise PyoframeError(
@@ -159,11 +297,24 @@ def _add_expressions_core(*expressions: "Expression") -> "Expression":
         propogate_strat = expressions[0].unmatched_strategy
         expr_data = [expr.data for expr in expressions]
 
+    # Add quadratic column if it is needed and doesn't already exist
+    if any(QUAD_VAR_KEY in df.columns for df in expr_data):
+        expr_data = [
+            (
+                df.with_columns(pl.lit(CONST_TERM).alias(QUAD_VAR_KEY).cast(KEY_TYPE))
+                if QUAD_VAR_KEY not in df.columns
+                else df
+            )
+            for df in expr_data
+        ]
+
     # Sort columns to allow for concat
-    expr_data = [e.select(sorted(e.columns)) for e in expr_data]
+    expr_data = [
+        e.select(dims + [c for c in e.columns if c not in dims]) for e in expr_data
+    ]
 
     data = pl.concat(expr_data, how="vertical_relaxed")
-    data = data.group_by(dims + [VAR_KEY], maintain_order=True).sum()
+    data = _sum_like_terms(data)
 
     new_expr = expressions[0]._new(data)
     new_expr.unmatched_strategy = propogate_strat
@@ -188,6 +339,7 @@ def _add_dimension(self: "Expression", target: "Expression") -> "Expression":
         return self
 
     if not set(missing_dims) <= set(self.allowed_new_dims):
+        # TODO actually suggest using e.g. .add_dim("a", "b") instead of just "use .add_dim()"
         raise PyoframeError(
             f"Dataframe has missing dimensions {missing_dims}. If this is intentional, use .add_dim()\n{self.data}"
         )
@@ -210,6 +362,20 @@ def _add_dimension(self: "Expression", target: "Expression") -> "Expression":
     return self._new(result)
 
 
+def _sum_like_terms(df: pl.DataFrame) -> pl.DataFrame:
+    """Combines terms with the same variables. Removes quadratic column if they all happen to cancel."""
+    dims = [c for c in df.columns if c not in RESERVED_COL_KEYS]
+    var_cols = [VAR_KEY] + ([QUAD_VAR_KEY] if QUAD_VAR_KEY in df.columns else [])
+    df = (
+        df.group_by(dims + var_cols, maintain_order=True)
+        .sum()
+        .filter(pl.col(COEF_KEY) != 0)
+    )
+    if QUAD_VAR_KEY in df.columns and (df.get_column(QUAD_VAR_KEY) == CONST_TERM).all():
+        df = df.drop(QUAD_VAR_KEY)
+    return df
+
+
 def _get_dimensions(df: pl.DataFrame) -> Optional[List[str]]:
     """
     Returns the dimensions of the DataFrame. Reserved columns do not count as dimensions.

pyoframe/constants.py

@@ -1,18 +1,14 @@
 """
 File containing shared constants used across the package.
-
-Code is heavily based on the `linopy` package by Fabian Hofmann.
-
-MIT License
 """
 
 import importlib.metadata
 import typing
-from dataclasses import dataclass
 from enum import Enum
-from typing import Literal, Optional, Union
+from typing import Literal, Optional
 
 import polars as pl
+import pyoptinterface as poi
 from packaging import version
 
 # We want to try and support multiple major versions of polars
@@ -20,22 +16,23 @@ POLARS_VERSION = version.parse(importlib.metadata.version("polars"))
 
 COEF_KEY = "__coeff"
 VAR_KEY = "__variable_id"
+QUAD_VAR_KEY = "__quadratic_variable_id"
 CONSTRAINT_KEY = "__constraint_id"
 SOLUTION_KEY = "solution"
 DUAL_KEY = "dual"
-RC_COL = "RC"
-SLACK_COL = "slack"
+SUPPORTED_SOLVER_TYPES = Literal["gurobi", "highs"]
+KEY_TYPE = pl.UInt32
 
+# Variable ID for constant terms. This variable ID is reserved.
 CONST_TERM = 0
 
 RESERVED_COL_KEYS = (
     COEF_KEY,
     VAR_KEY,
+    QUAD_VAR_KEY,
     CONSTRAINT_KEY,
     SOLUTION_KEY,
     DUAL_KEY,
-    RC_COL,
-    SLACK_COL,
 )
 
 
@@ -52,9 +49,16 @@ class _ConfigMeta(type):
 
 
 class Config(metaclass=_ConfigMeta):
+    """
+    Configuration options that apply to the entire library.
+    """
+
+    default_solver: Optional[SUPPORTED_SOLVER_TYPES] = None
     disable_unmatched_checks: bool = False
-    print_float_precision: Optional[int] = 5
+    float_to_str_precision: Optional[int] = 5
     print_uses_variable_names: bool = True
+    print_max_line_length: int = 80
+    print_max_lines: int = 15
     # Number of elements to show when printing a set to the console (additional elements are replaced with ...)
     print_max_set_elements: int = 50
     enable_is_duplicated_expression_safety_check: bool = False
@@ -73,17 +77,45 @@ class ConstraintSense(Enum):
     GE = ">="
     EQ = "="
 
+    def to_poi(self):
+        if self == ConstraintSense.LE:
+            return poi.ConstraintSense.LessEqual
+        elif self == ConstraintSense.EQ:
+            return poi.ConstraintSense.Equal
+        elif self == ConstraintSense.GE:
+            return poi.ConstraintSense.GreaterEqual
+        else:
+            raise ValueError(f"Invalid constraint type: {self}")  # pragma: no cover
+
 
 class ObjSense(Enum):
     MIN = "min"
     MAX = "max"
 
+    def to_poi(self):
+        if self == ObjSense.MIN:
+            return poi.ObjectiveSense.Minimize
+        elif self == ObjSense.MAX:
+            return poi.ObjectiveSense.Maximize
+        else:
+            raise ValueError(f"Invalid objective sense: {self}")  # pragma: no cover
+
 
 class VType(Enum):
     CONTINUOUS = "continuous"
     BINARY = "binary"
     INTEGER = "integer"
 
+    def to_poi(self):
+        if self == VType.CONTINUOUS:
+            return poi.VariableDomain.Continuous
+        elif self == VType.BINARY:
+            return poi.VariableDomain.Binary
+        elif self == VType.INTEGER:
+            return poi.VariableDomain.Integer
+        else:
+            raise ValueError(f"Invalid variable type: {self}")  # pragma: no cover
+
 
 class UnmatchedStrategy(Enum):
     UNSET = "not_set"
@@ -99,193 +131,5 @@ 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,
-            termination_conditions,
-        ) in STATUS_TO_TERMINATION_CONDITION_MAP.items():
-            if termination_condition in termination_conditions:
-                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
-
-    def __repr__(self) -> str:
-        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"
-            )
-
-        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")
-
-
 class PyoframeError(Exception):
     pass