Index

Pyoframe's public API. Also applies the monkey patch to the DataFrame libraries.

Config

Configuration options that apply to the entire library.

reset_defaults() classmethod

Resets all configuration options to their default values.

Source code in pyoframe/constants.py

@classmethod
def reset_defaults(cls):
    """
    Resets all configuration options to their default values.
    """
    for key, value in cls._defaults.items():
        setattr(cls, key, value)

Constraint(lhs, sense)

Bases: ModelElementWithId

A linear programming constraint.

Parameters:

Name	Type	Description	Default
lhs	Expression	The left hand side of the constraint.	required
sense	ConstraintSense	The sense of the constraint.	required

Source code in pyoframe/core.py

def __init__(self, lhs: Expression, sense: ConstraintSense):
    """Initialize a constraint.

    Parameters:
        lhs:
            The left hand side of the constraint.
        sense:
            The sense of the constraint.
    """
    self.lhs = lhs
    self._model = lhs._model
    self.sense = sense
    self.to_relax: Optional[FuncArgs] = None
    self.attr = Container(self._set_attribute, self._get_attribute)

    dims = self.lhs.dimensions
    data = pl.DataFrame() if dims is None else self.lhs.data.select(dims).unique()

    super().__init__(data)

relax(cost, max=None)

Relaxes the constraint by adding a variable to the constraint that can be non-zero at a cost.

Parameters:

Name	Type	Description	Default
cost	SupportsToExpr	The cost of relaxing the constraint. Costs should be positives as they will automatically become negative for maximization problems.	required
max	Optional[SupportsToExpr]	The maximum value of the relaxation variable.	None

Returns:

Type	Description
Constraint	The same constraint

Examples:

>>> m = pf.Model()
>>> m.hours_sleep = pf.Variable(lb=0)
>>> m.hours_day = pf.Variable(lb=0)
>>> m.hours_in_day = m.hours_sleep + m.hours_day == 24
>>> m.maximize = m.hours_day
>>> m.must_sleep = (m.hours_sleep >= 8).relax(cost=2, max=3)
>>> m.optimize()
>>> m.hours_day.solution
16.0
>>> m.maximize += 2 * m.hours_day
>>> m.optimize()
>>> m.hours_day.solution
19.0

Note: .relax() can only be called after the sense of the model has been defined.

>>> m = pf.Model()
>>> m.hours_sleep = pf.Variable(lb=0)
>>> m.hours_day = pf.Variable(lb=0)
>>> m.hours_in_day = m.hours_sleep + m.hours_day == 24
>>> m.must_sleep = (m.hours_sleep >= 8).relax(cost=2, max=3)
Traceback (most recent call last):
...
ValueError: Cannot relax a constraint before the objective sense has been set. Try setting the objective first or using Model(sense=...).

One way to solve this is by setting the sense directly on the model. See how this works fine:

>>> m = pf.Model(sense="max")
>>> m.hours_sleep = pf.Variable(lb=0)
>>> m.hours_day = pf.Variable(lb=0)
>>> m.hours_in_day = m.hours_sleep + m.hours_day == 24
>>> m.must_sleep = (m.hours_sleep >= 8).relax(cost=2, max=3)

And now an example with dimensions:

>>> homework_due_tomorrow = pl.DataFrame({"project": ["A", "B", "C"], "cost_per_hour_underdelivered": [10, 20, 30], "hours_to_finish": [9, 9, 9], "max_underdelivered": [1, 9, 9]})
>>> m.hours_spent = pf.Variable(homework_due_tomorrow[["project"]], lb=0)
>>> m.must_finish_project = (m.hours_spent >= homework_due_tomorrow[["project", "hours_to_finish"]]).relax(homework_due_tomorrow[["project", "cost_per_hour_underdelivered"]], max=homework_due_tomorrow[["project", "max_underdelivered"]])
>>> m.only_one_day = sum("project", m.hours_spent) <= 24
>>> # Relaxing a constraint after it has already been assigned will give an error
>>> m.only_one_day.relax(1)
Traceback (most recent call last):
...
ValueError: .relax() must be called before the Constraint is added to the model
>>> m.attr.Silent = True
>>> m.optimize()
>>> m.maximize.value
-50.0
>>> m.hours_spent.solution
shape: (3, 2)
┌─────────┬──────────┐
│ project ┆ solution │
│ ---     ┆ ---      │
│ str     ┆ f64      │
╞═════════╪══════════╡
│ A       ┆ 8.0      │
│ B       ┆ 7.0      │
│ C       ┆ 9.0      │
└─────────┴──────────┘

Source code in pyoframe/core.py

def relax(
    self, cost: SupportsToExpr, max: Optional[SupportsToExpr] = None
) -> Constraint:
    """
    Relaxes the constraint by adding a variable to the constraint that can be non-zero at a cost.

    Parameters:
        cost:
            The cost of relaxing the constraint. Costs should be positives as they will automatically
            become negative for maximization problems.
        max:
            The maximum value of the relaxation variable.

    Returns:
        The same constraint

    Examples:
        >>> m = pf.Model()
        >>> m.hours_sleep = pf.Variable(lb=0)
        >>> m.hours_day = pf.Variable(lb=0)
        >>> m.hours_in_day = m.hours_sleep + m.hours_day == 24
        >>> m.maximize = m.hours_day
        >>> m.must_sleep = (m.hours_sleep >= 8).relax(cost=2, max=3)
        >>> m.optimize()
        >>> m.hours_day.solution
        16.0
        >>> m.maximize += 2 * m.hours_day
        >>> m.optimize()
        >>> m.hours_day.solution
        19.0

        Note: .relax() can only be called after the sense of the model has been defined.

        >>> m = pf.Model()
        >>> m.hours_sleep = pf.Variable(lb=0)
        >>> m.hours_day = pf.Variable(lb=0)
        >>> m.hours_in_day = m.hours_sleep + m.hours_day == 24
        >>> m.must_sleep = (m.hours_sleep >= 8).relax(cost=2, max=3)
        Traceback (most recent call last):
        ...
        ValueError: Cannot relax a constraint before the objective sense has been set. Try setting the objective first or using Model(sense=...).

        One way to solve this is by setting the sense directly on the model. See how this works fine:

        >>> m = pf.Model(sense="max")
        >>> m.hours_sleep = pf.Variable(lb=0)
        >>> m.hours_day = pf.Variable(lb=0)
        >>> m.hours_in_day = m.hours_sleep + m.hours_day == 24
        >>> m.must_sleep = (m.hours_sleep >= 8).relax(cost=2, max=3)

        And now an example with dimensions:

        >>> homework_due_tomorrow = pl.DataFrame({"project": ["A", "B", "C"], "cost_per_hour_underdelivered": [10, 20, 30], "hours_to_finish": [9, 9, 9], "max_underdelivered": [1, 9, 9]})
        >>> m.hours_spent = pf.Variable(homework_due_tomorrow[["project"]], lb=0)
        >>> m.must_finish_project = (m.hours_spent >= homework_due_tomorrow[["project", "hours_to_finish"]]).relax(homework_due_tomorrow[["project", "cost_per_hour_underdelivered"]], max=homework_due_tomorrow[["project", "max_underdelivered"]])
        >>> m.only_one_day = sum("project", m.hours_spent) <= 24
        >>> # Relaxing a constraint after it has already been assigned will give an error
        >>> m.only_one_day.relax(1)
        Traceback (most recent call last):
        ...
        ValueError: .relax() must be called before the Constraint is added to the model
        >>> m.attr.Silent = True
        >>> m.optimize()
        >>> m.maximize.value
        -50.0
        >>> m.hours_spent.solution
        shape: (3, 2)
        ┌─────────┬──────────┐
        │ project ┆ solution │
        │ ---     ┆ ---      │
        │ str     ┆ f64      │
        ╞═════════╪══════════╡
        │ A       ┆ 8.0      │
        │ B       ┆ 7.0      │
        │ C       ┆ 9.0      │
        └─────────┴──────────┘
    """
    if self._has_ids:
        raise ValueError(
            ".relax() must be called before the Constraint is added to the model"
        )

    m = self._model
    if m is None or self.name is None:
        self.to_relax = FuncArgs(args=[cost, max])
        return self

    var_name = f"{self.name}_relaxation"
    assert not hasattr(
        m, var_name
    ), "Conflicting names, relaxation variable already exists on the model."
    var = Variable(self, lb=0, ub=max)
    setattr(m, var_name, var)

    if self.sense == ConstraintSense.LE:
        self.lhs -= var
    elif self.sense == ConstraintSense.GE:
        self.lhs += var
    else:  # pragma: no cover
        # TODO
        raise NotImplementedError(
            "Relaxation for equalities has not yet been implemented. Submit a pull request!"
        )

    penalty = var * cost
    if self.dimensions:
        penalty = sum(self.dimensions, penalty)
    if m.sense is None:
        raise ValueError(
            "Cannot relax a constraint before the objective sense has been set. Try setting the objective first or using Model(sense=...)."
        )
    elif m.sense == ObjSense.MAX:
        penalty *= -1
    if m.objective is None:
        m.objective = penalty
    else:
        m.objective += penalty

    return self

Expression(data)

Bases: ModelElement, SupportsMath, SupportPolarsMethodMixin

A linear or quadratic expression.

Examples:

>>> import pandas as pd
>>> df = pd.DataFrame({"item" : [1, 1, 1, 2, 2], "time": ["mon", "tue", "wed", "mon", "tue"], "cost": [1, 2, 3, 4, 5]}).set_index(["item", "time"])
>>> m = pf.Model()
>>> m.Time = pf.Variable(df.index)
>>> m.Size = pf.Variable(df.index)
>>> expr = df["cost"] * m.Time + df["cost"] * m.Size
>>> expr
<Expression size=5 dimensions={'item': 2, 'time': 3} terms=10>
[1,mon]: Time[1,mon] + Size[1,mon]
[1,tue]: 2 Time[1,tue] +2 Size[1,tue]
[1,wed]: 3 Time[1,wed] +3 Size[1,wed]
[2,mon]: 4 Time[2,mon] +4 Size[2,mon]
[2,tue]: 5 Time[2,tue] +5 Size[2,tue]

Source code in pyoframe/core.py

def __init__(self, data: pl.DataFrame):
    """
    A linear expression.

    Examples:
        >>> import pandas as pd
        >>> df = pd.DataFrame({"item" : [1, 1, 1, 2, 2], "time": ["mon", "tue", "wed", "mon", "tue"], "cost": [1, 2, 3, 4, 5]}).set_index(["item", "time"])
        >>> m = pf.Model()
        >>> m.Time = pf.Variable(df.index)
        >>> m.Size = pf.Variable(df.index)
        >>> expr = df["cost"] * m.Time + df["cost"] * m.Size
        >>> expr
        <Expression size=5 dimensions={'item': 2, 'time': 3} terms=10>
        [1,mon]: Time[1,mon] + Size[1,mon]
        [1,tue]: 2 Time[1,tue] +2 Size[1,tue]
        [1,wed]: 3 Time[1,wed] +3 Size[1,wed]
        [2,mon]: 4 Time[2,mon] +4 Size[2,mon]
        [2,tue]: 5 Time[2,tue] +5 Size[2,tue]
    """
    # Sanity checks, VAR_KEY and COEF_KEY must be present
    assert VAR_KEY in data.columns, "Missing variable column."
    assert COEF_KEY in data.columns, "Missing coefficient column."

    # Sanity check no duplicates indices
    if Config.enable_is_duplicated_expression_safety_check:
        duplicated_mask = data.drop(COEF_KEY).is_duplicated()
        # In theory this should never happen unless there's a bug in the library
        if duplicated_mask.any():  # pragma: no cover
            duplicated_data = data.filter(duplicated_mask)
            raise ValueError(
                f"Cannot create an expression with duplicate indices:\n{duplicated_data}."
            )

    data = _simplify_expr_df(data)

    super().__init__(data)

is_quadratic property

Returns True if the expression is quadratic, False otherwise.

Computes in O(1) since expressions are quadratic if and only if self.data contain the QUAD_VAR_KEY column.

Examples:

>>> import pandas as pd
>>> m = pf.Model()
>>> m.v = Variable()
>>> expr = pd.DataFrame({"dim1": [1, 2, 3], "value": [1, 2, 3]}) * m.v
>>> expr *= m.v
>>> expr.is_quadratic
True

terms property

Number of terms across all subexpressions.

Expressions equal to zero count as one term.

Examples:

>>> import polars as pl
>>> m = pf.Model()
>>> m.v = pf.Variable({"t": [1, 2]})
>>> coef = pl.DataFrame({"t": [1, 2], "coef": [0, 1]})
>>> coef*(m.v+4)
<Expression size=2 dimensions={'t': 2} terms=3>
[1]: 0
[2]: 4  + v[2]
>>> (coef*(m.v+4)).terms
3

__add__(other)

Examples:

>>> import pandas as pd
>>> m = pf.Model()
>>> add = pd.DataFrame({"dim1": [1,2,3], "add": [10, 20, 30]}).to_expr()
>>> m.v = Variable(add)
>>> m.v + add
<Expression size=3 dimensions={'dim1': 3} terms=6>
[1]: v[1] +10
[2]: v[2] +20
[3]: v[3] +30
>>> m.v + add + 2
<Expression size=3 dimensions={'dim1': 3} terms=6>
[1]: v[1] +12
[2]: v[2] +22
[3]: v[3] +32
>>> m.v + pd.DataFrame({"dim1": [1,2], "add": [10, 20]})
Traceback (most recent call last):
...
pyoframe.constants.PyoframeError: Failed to add expressions:
<Expression size=3 dimensions={'dim1': 3} terms=3> + <Expression size=2 dimensions={'dim1': 2} terms=2>
Due to error:
Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()
shape: (1, 2)
┌──────┬────────────┐
│ dim1 ┆ dim1_right │
│ ---  ┆ ---        │
│ i64  ┆ i64        │
╞══════╪════════════╡
│ 3    ┆ null       │
└──────┴────────────┘
>>> m.v2 = Variable()
>>> 5 + 2 * m.v2
<Expression size=1 dimensions={} terms=2>
2 v2 +5

Source code in pyoframe/core.py

def __add__(self, other):
    """
    Examples:
        >>> import pandas as pd
        >>> m = pf.Model()
        >>> add = pd.DataFrame({"dim1": [1,2,3], "add": [10, 20, 30]}).to_expr()
        >>> m.v = Variable(add)
        >>> m.v + add
        <Expression size=3 dimensions={'dim1': 3} terms=6>
        [1]: v[1] +10
        [2]: v[2] +20
        [3]: v[3] +30
        >>> m.v + add + 2
        <Expression size=3 dimensions={'dim1': 3} terms=6>
        [1]: v[1] +12
        [2]: v[2] +22
        [3]: v[3] +32
        >>> m.v + pd.DataFrame({"dim1": [1,2], "add": [10, 20]})
        Traceback (most recent call last):
        ...
        pyoframe.constants.PyoframeError: Failed to add expressions:
        <Expression size=3 dimensions={'dim1': 3} terms=3> + <Expression size=2 dimensions={'dim1': 2} terms=2>
        Due to error:
        Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()
        shape: (1, 2)
        ┌──────┬────────────┐
        │ dim1 ┆ dim1_right │
        │ ---  ┆ ---        │
        │ i64  ┆ i64        │
        ╞══════╪════════════╡
        │ 3    ┆ null       │
        └──────┴────────────┘
        >>> m.v2 = Variable()
        >>> 5 + 2 * m.v2
        <Expression size=1 dimensions={} terms=2>
        2 v2 +5
    """
    if isinstance(other, str):
        raise ValueError(
            "Cannot add a string to an expression. Perhaps you meant to use pf.sum() instead of sum()?"
        )
    if isinstance(other, (int, float)):
        return self._add_const(other)
    other = other.to_expr()
    self._learn_from_other(other)
    return _add_expressions(self, other)

_add_const(const)

Examples:

>>> m = pf.Model()
>>> m.x1 = Variable()
>>> m.x2 = Variable()
>>> m.x1 + 5
<Expression size=1 dimensions={} terms=2>
x1 +5
>>> m.x1 ** 2 + 5
<Expression size=1 dimensions={} terms=2 degree=2>
x1 * x1 +5
>>> m.x1 ** 2 + m.x2 + 5
<Expression size=1 dimensions={} terms=3 degree=2>
x1 * x1 + x2 +5

It also works with dimensions

>>> m = pf.Model()
>>> m.v = Variable({"dim1": [1, 2, 3]})
>>> m.v * m.v + 5
<Expression size=3 dimensions={'dim1': 3} terms=6 degree=2>
[1]: 5  + v[1] * v[1]
[2]: 5  + v[2] * v[2]
[3]: 5  + v[3] * v[3]

Source code in pyoframe/core.py

def _add_const(self, const: int | float) -> Expression:
    """
    Examples:
        >>> m = pf.Model()
        >>> m.x1 = Variable()
        >>> m.x2 = Variable()
        >>> m.x1 + 5
        <Expression size=1 dimensions={} terms=2>
        x1 +5
        >>> m.x1 ** 2 + 5
        <Expression size=1 dimensions={} terms=2 degree=2>
        x1 * x1 +5
        >>> m.x1 ** 2 + m.x2 + 5
        <Expression size=1 dimensions={} terms=3 degree=2>
        x1 * x1 + x2 +5

        It also works with dimensions

        >>> m = pf.Model()
        >>> m.v = Variable({"dim1": [1, 2, 3]})
        >>> m.v * m.v + 5
        <Expression size=3 dimensions={'dim1': 3} terms=6 degree=2>
        [1]: 5  + v[1] * v[1]
        [2]: 5  + v[2] * v[2]
        [3]: 5  + v[3] * v[3]
    """
    dim = self.dimensions
    data = self.data
    # Fill in missing constant terms
    if not dim:
        if CONST_TERM not in data[VAR_KEY]:
            const_df = pl.DataFrame(
                {COEF_KEY: [0.0], VAR_KEY: [CONST_TERM]},
                schema={COEF_KEY: pl.Float64, VAR_KEY: KEY_TYPE},
            )
            if self.is_quadratic:
                const_df = const_df.with_columns(
                    pl.lit(CONST_TERM).alias(QUAD_VAR_KEY).cast(KEY_TYPE)
                )
            data = pl.concat(
                [data, const_df],
                how="vertical_relaxed",
            )
    else:
        keys = (
            data.select(dim)
            .unique(maintain_order=True)
            .with_columns(pl.lit(CONST_TERM).alias(VAR_KEY).cast(KEY_TYPE))
        )
        if self.is_quadratic:
            keys = keys.with_columns(
                pl.lit(CONST_TERM).alias(QUAD_VAR_KEY).cast(KEY_TYPE)
            )
        if POLARS_VERSION.major >= 1:
            data = data.join(
                keys, on=dim + self._variable_columns, how="full", coalesce=True
            )
        else:
            data = data.join(
                keys, on=dim + self._variable_columns, how="outer_coalesce"
            )
        data = data.with_columns(pl.col(COEF_KEY).fill_null(0.0))

    data = data.with_columns(
        pl.when(pl.col(VAR_KEY) == CONST_TERM)
        .then(pl.col(COEF_KEY) + const)
        .otherwise(pl.col(COEF_KEY))
    )

    return self._new(data)

constant(constant) classmethod

Examples:

>>> pf.Expression.constant(5)
<Expression size=1 dimensions={} terms=1>
5

Source code in pyoframe/core.py

@classmethod
def constant(cls, constant: int | float) -> "Expression":
    """
    Examples:
        >>> pf.Expression.constant(5)
        <Expression size=1 dimensions={} terms=1>
        5
    """
    return cls(
        pl.DataFrame(
            {
                COEF_KEY: [constant],
                VAR_KEY: [CONST_TERM],
            },
            schema={COEF_KEY: pl.Float64, VAR_KEY: KEY_TYPE},
        )
    )

degree()

Returns the degree of the expression (0=constant, 1=linear, 2=quadratic).

Examples:

>>> import pandas as pd
>>> m = pf.Model()
>>> m.v1 = pf.Variable()
>>> m.v2 = pf.Variable()
>>> expr = pd.DataFrame({"dim1": [1, 2, 3], "value": [1, 2, 3]}).to_expr()
>>> expr.degree()
0
>>> expr *= m.v1
>>> expr.degree()
1
>>> expr += (m.v2 ** 2).add_dim("dim1")
>>> expr.degree()
2

Source code in pyoframe/core.py

def degree(self) -> int:
    """
    Returns the degree of the expression (0=constant, 1=linear, 2=quadratic).

    Examples:
        >>> import pandas as pd
        >>> m = pf.Model()
        >>> m.v1 = pf.Variable()
        >>> m.v2 = pf.Variable()
        >>> expr = pd.DataFrame({"dim1": [1, 2, 3], "value": [1, 2, 3]}).to_expr()
        >>> expr.degree()
        0
        >>> expr *= m.v1
        >>> expr.degree()
        1
        >>> expr += (m.v2 ** 2).add_dim("dim1")
        >>> expr.degree()
        2
    """
    if self.is_quadratic:
        return 2
    elif (self.data.get_column(VAR_KEY) != CONST_TERM).any():
        return 1
    else:
        return 0

evaluate()

The value of the expression. Only available after the model has been solved.

Examples:

>>> m = pf.Model()
>>> m.X = pf.Variable({"dim1": [1, 2, 3]}, ub=10)
>>> m.expr_1 = 2 * m.X + 1
>>> m.expr_2 = pf.sum(m.expr_1)
>>> m.maximize = m.expr_2 - 3
>>> m.attr.Silent = True
>>> m.optimize()
>>> m.expr_1.evaluate()
shape: (3, 2)
┌──────┬──────────┐
│ dim1 ┆ solution │
│ ---  ┆ ---      │
│ i64  ┆ f64      │
╞══════╪══════════╡
│ 1    ┆ 21.0     │
│ 2    ┆ 21.0     │
│ 3    ┆ 21.0     │
└──────┴──────────┘
>>> m.expr_2.evaluate()
63.0

Source code in pyoframe/core.py

@unwrap_single_values
def evaluate(self) -> pl.DataFrame:
    """
    The value of the expression. Only available after the model has been solved.

    Examples:
        >>> m = pf.Model()
        >>> m.X = pf.Variable({"dim1": [1, 2, 3]}, ub=10)
        >>> m.expr_1 = 2 * m.X + 1
        >>> m.expr_2 = pf.sum(m.expr_1)
        >>> m.maximize = m.expr_2 - 3
        >>> m.attr.Silent = True
        >>> m.optimize()
        >>> m.expr_1.evaluate()
        shape: (3, 2)
        ┌──────┬──────────┐
        │ dim1 ┆ solution │
        │ ---  ┆ ---      │
        │ i64  ┆ f64      │
        ╞══════╪══════════╡
        │ 1    ┆ 21.0     │
        │ 2    ┆ 21.0     │
        │ 3    ┆ 21.0     │
        └──────┴──────────┘
        >>> m.expr_2.evaluate()
        63.0
    """
    assert (
        self._model is not None
    ), "Expression must be added to the model to use .value"

    df = self.data
    sm = self._model.poi
    attr = poi.VariableAttribute.Value
    for var_col in self._variable_columns:
        df = df.with_columns(
            (
                pl.col(COEF_KEY)
                * pl.col(var_col).map_elements(
                    lambda v_id: (
                        sm.get_variable_attribute(poi.VariableIndex(v_id), attr)
                        if v_id != CONST_TERM
                        else 1
                    ),
                    return_dtype=pl.Float64,
                )
            ).alias(COEF_KEY)
        ).drop(var_col)

    df = df.rename({COEF_KEY: SOLUTION_KEY})

    dims = self.dimensions
    if dims is not None:
        df = df.group_by(dims, maintain_order=True)
    return df.sum()

map(mapping_set, drop_shared_dims=True)

Replaces the dimensions that are shared with mapping_set with the other dimensions found in mapping_set.

This is particularly useful to go from one type of dimensions to another. For example, to convert data that is indexed by city to data indexed by country (see example).

Parameters:

Name	Type	Description	Default
mapping_set	SetTypes	The set to map the expression to. This can be a DataFrame, Index, or another Set.	required
drop_shared_dims	bool	If True, the dimensions shared between the expression and the mapping set are dropped from the resulting expression and repeated rows are summed. If False, the shared dimensions are kept in the resulting expression.	True

Returns:

Type	Description
Expression	A new Expression containing the result of the mapping operation.

Examples:

>>> import polars as pl
>>> pop_data = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "year": [2024, 2024, 2024], "population": [10, 2, 8]}).to_expr()
>>> cities_and_countries = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "country": ["Canada", "Canada", "USA"]})
>>> pop_data.map(cities_and_countries)
<Expression size=2 dimensions={'year': 1, 'country': 2} terms=2>
[2024,Canada]: 12
[2024,USA]: 8
>>> pop_data.map(cities_and_countries, drop_shared_dims=False)
<Expression size=3 dimensions={'city': 3, 'year': 1, 'country': 2} terms=3>
[Toronto,2024,Canada]: 10
[Vancouver,2024,Canada]: 2
[Boston,2024,USA]: 8

Source code in pyoframe/core.py

def map(self, mapping_set: SetTypes, drop_shared_dims: bool = True) -> Expression:
    """
    Replaces the dimensions that are shared with mapping_set with the other dimensions found in mapping_set.

    This is particularly useful to go from one type of dimensions to another. For example, to convert data that
    is indexed by city to data indexed by country (see example).

    Parameters:
        mapping_set:
            The set to map the expression to. This can be a DataFrame, Index, or another Set.
        drop_shared_dims:
            If True, the dimensions shared between the expression and the mapping set are dropped from the resulting expression and
                repeated rows are summed.
            If False, the shared dimensions are kept in the resulting expression.

    Returns:
        A new Expression containing the result of the mapping operation.

    Examples:
        >>> import polars as pl
        >>> pop_data = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "year": [2024, 2024, 2024], "population": [10, 2, 8]}).to_expr()
        >>> cities_and_countries = pl.DataFrame({"city": ["Toronto", "Vancouver", "Boston"], "country": ["Canada", "Canada", "USA"]})
        >>> pop_data.map(cities_and_countries)
        <Expression size=2 dimensions={'year': 1, 'country': 2} terms=2>
        [2024,Canada]: 12
        [2024,USA]: 8
        >>> pop_data.map(cities_and_countries, drop_shared_dims=False)
        <Expression size=3 dimensions={'city': 3, 'year': 1, 'country': 2} terms=3>
        [Toronto,2024,Canada]: 10
        [Vancouver,2024,Canada]: 2
        [Boston,2024,USA]: 8
    """
    mapping_set = Set(mapping_set)

    dims = self.dimensions
    if dims is None:
        raise ValueError("Cannot use .map() on an expression with no dimensions.")

    mapping_dims = mapping_set.dimensions
    if mapping_dims is None:
        raise ValueError(
            "Cannot use .map() with a mapping set containing no dimensions."
        )

    shared_dims = [dim for dim in dims if dim in mapping_dims]
    if not shared_dims:
        raise ValueError(
            f"Cannot apply .map() as there are no shared dimensions between the expression (dims={self.dimensions}) and the mapping set (dims={mapping_set.dimensions})."
        )

    mapped_expression = self * mapping_set

    if drop_shared_dims:
        return sum(shared_dims, mapped_expression)

    return mapped_expression

rolling_sum(over, window_size)

Calculates the rolling sum of the Expression over a specified window size for a given dimension.

This method applies a rolling sum operation over the dimension specified by over, using a window defined by window_size.

Parameters:

Name	Type	Description	Default
over		The name of the dimension (column) over which the rolling sum is calculated. This dimension must exist within the Expression's dimensions.	required
window_size		The size of the moving window in terms of number of records. The rolling sum is calculated over this many consecutive elements.	required

Returns:

Type	Description
Expression	A new Expression instance containing the result of the rolling sum operation. This new Expression retains all dimensions (columns) of the original data, with the rolling sum applied over the specified dimension.

Examples:

>>> import polars as pl
>>> cost = pl.DataFrame({"item" : [1, 1, 1, 2, 2], "time": [1, 2, 3, 1, 2], "cost": [1, 2, 3, 4, 5]})
>>> m = pf.Model()
>>> m.quantity = pf.Variable(cost[["item", "time"]])
>>> (m.quantity * cost).rolling_sum(over="time", window_size=2)
<Expression size=5 dimensions={'item': 2, 'time': 3} terms=8>
[1,1]: quantity[1,1]
[1,2]: quantity[1,1] +2 quantity[1,2]
[1,3]: 2 quantity[1,2] +3 quantity[1,3]
[2,1]: 4 quantity[2,1]
[2,2]: 4 quantity[2,1] +5 quantity[2,2]

Source code in pyoframe/core.py

def rolling_sum(self, over: str, window_size: int) -> Expression:
    """
    Calculates the rolling sum of the Expression over a specified window size for a given dimension.

    This method applies a rolling sum operation over the dimension specified by `over`,
    using a window defined by `window_size`.


    Parameters:
        over :
            The name of the dimension (column) over which the rolling sum is calculated.
            This dimension must exist within the Expression's dimensions.
        window_size :
            The size of the moving window in terms of number of records.
            The rolling sum is calculated over this many consecutive elements.

    Returns:
        A new Expression instance containing the result of the rolling sum operation.
            This new Expression retains all dimensions (columns) of the original data,
            with the rolling sum applied over the specified dimension.

    Examples:
        >>> import polars as pl
        >>> cost = pl.DataFrame({"item" : [1, 1, 1, 2, 2], "time": [1, 2, 3, 1, 2], "cost": [1, 2, 3, 4, 5]})
        >>> m = pf.Model()
        >>> m.quantity = pf.Variable(cost[["item", "time"]])
        >>> (m.quantity * cost).rolling_sum(over="time", window_size=2)
        <Expression size=5 dimensions={'item': 2, 'time': 3} terms=8>
        [1,1]: quantity[1,1]
        [1,2]: quantity[1,1] +2 quantity[1,2]
        [1,3]: 2 quantity[1,2] +3 quantity[1,3]
        [2,1]: 4 quantity[2,1]
        [2,2]: 4 quantity[2,1] +5 quantity[2,2]
    """
    dims = self.dimensions
    if dims is None:
        raise ValueError(
            "Cannot use rolling_sum() with an expression with no dimensions."
        )
    assert over in dims, f"Cannot sum over {over} as it is not in {dims}"
    remaining_dims = [dim for dim in dims if dim not in over]

    return self._new(
        pl.concat(
            [
                df.with_columns(pl.col(over).max())
                for _, df in self.data.rolling(
                    index_column=over,
                    period=f"{window_size}i",
                    group_by=remaining_dims,
                )
            ]
        )
    )

sum(over)

Examples:

>>> import pandas as pd
>>> m = pf.Model()
>>> df = pd.DataFrame({"item" : [1, 1, 1, 2, 2], "time": ["mon", "tue", "wed", "mon", "tue"], "cost": [1, 2, 3, 4, 5]}).set_index(["item", "time"])
>>> m.quantity = Variable(df.reset_index()[["item"]].drop_duplicates())
>>> expr = (m.quantity * df["cost"]).sum("time")
>>> expr.data
shape: (2, 3)
┌──────┬─────────┬───────────────┐
│ item ┆ __coeff ┆ __variable_id │
│ ---  ┆ ---     ┆ ---           │
│ i64  ┆ f64     ┆ u32           │
╞══════╪═════════╪═══════════════╡
│ 1    ┆ 6.0     ┆ 1             │
│ 2    ┆ 9.0     ┆ 2             │
└──────┴─────────┴───────────────┘

Source code in pyoframe/core.py

def sum(self, over: Union[str, Iterable[str]]):
    """
    Examples:
        >>> import pandas as pd
        >>> m = pf.Model()
        >>> df = pd.DataFrame({"item" : [1, 1, 1, 2, 2], "time": ["mon", "tue", "wed", "mon", "tue"], "cost": [1, 2, 3, 4, 5]}).set_index(["item", "time"])
        >>> m.quantity = Variable(df.reset_index()[["item"]].drop_duplicates())
        >>> expr = (m.quantity * df["cost"]).sum("time")
        >>> expr.data
        shape: (2, 3)
        ┌──────┬─────────┬───────────────┐
        │ item ┆ __coeff ┆ __variable_id │
        │ ---  ┆ ---     ┆ ---           │
        │ i64  ┆ f64     ┆ u32           │
        ╞══════╪═════════╪═══════════════╡
        │ 1    ┆ 6.0     ┆ 1             │
        │ 2    ┆ 9.0     ┆ 2             │
        └──────┴─────────┴───────────────┘
    """
    if isinstance(over, str):
        over = [over]
    dims = self.dimensions
    if not dims:
        raise ValueError(
            f"Cannot sum over dimensions {over} since the current expression has no dimensions."
        )
    assert set(over) <= set(dims), f"Cannot sum over {over} as it is not in {dims}"
    remaining_dims = [dim for dim in dims if dim not in over]

    return self._new(
        self.data.drop(over)
        .group_by(remaining_dims + self._variable_columns, maintain_order=True)
        .sum()
    )

within(set)

Examples:

>>> import pandas as pd
>>> general_expr = pd.DataFrame({"dim1": [1, 2, 3], "value": [1, 2, 3]}).to_expr()
>>> filter_expr = pd.DataFrame({"dim1": [1, 3], "value": [5, 6]}).to_expr()
>>> general_expr.within(filter_expr).data
shape: (2, 3)
┌──────┬─────────┬───────────────┐
│ dim1 ┆ __coeff ┆ __variable_id │
│ ---  ┆ ---     ┆ ---           │
│ i64  ┆ f64     ┆ u32           │
╞══════╪═════════╪═══════════════╡
│ 1    ┆ 1.0     ┆ 0             │
│ 3    ┆ 3.0     ┆ 0             │
└──────┴─────────┴───────────────┘

Source code in pyoframe/core.py

def within(self, set: "SetTypes") -> Expression:
    """
    Examples:
        >>> import pandas as pd
        >>> general_expr = pd.DataFrame({"dim1": [1, 2, 3], "value": [1, 2, 3]}).to_expr()
        >>> filter_expr = pd.DataFrame({"dim1": [1, 3], "value": [5, 6]}).to_expr()
        >>> general_expr.within(filter_expr).data
        shape: (2, 3)
        ┌──────┬─────────┬───────────────┐
        │ dim1 ┆ __coeff ┆ __variable_id │
        │ ---  ┆ ---     ┆ ---           │
        │ i64  ┆ f64     ┆ u32           │
        ╞══════╪═════════╪═══════════════╡
        │ 1    ┆ 1.0     ┆ 0             │
        │ 3    ┆ 3.0     ┆ 0             │
        └──────┴─────────┴───────────────┘
    """
    df: pl.DataFrame = Set(set).data
    set_dims = _get_dimensions(df)
    assert (
        set_dims is not None
    ), "Cannot use .within() with a set with no dimensions."
    dims = self.dimensions
    assert (
        dims is not None
    ), "Cannot use .within() with an expression with no dimensions."
    dims_in_common = [dim for dim in dims if dim in set_dims]
    by_dims = df.select(dims_in_common).unique(maintain_order=True)
    return self._new(self.data.join(by_dims, on=dims_in_common))

Model(name=None, solver=None, solver_env=None, use_var_names=False, sense=None)

The object that holds all the variables, constraints, and the objective.

Parameters:

Name	Type	Description	Default
name		The name of the model. Currently it is not used for much.	None
solver	Optional[SUPPORTED_SOLVER_TYPES]	The solver to use. If None, Config.default_solver will be used. If Config.default_solver has not been set (None), Pyoframe will try to detect whichever solver is already installed.	None
solver_env	Optional[Dict[str, str]]	Gurobi only: a dictionary of parameters to set when creating the Gurobi environment.	None
use_var_names		Whether to pass variable names to the solver. Set to True if you'd like outputs from e.g. Model.write() to be legible.	False
sense	Union[ObjSense, ObjSenseValue, None]	Either "min" or "max". Indicates whether it's a minmization or maximization problem. Typically, this parameter can be omitted (None) as it will automatically be set when the objective is set using .minimize or .maximize.	None

Examples:

>>> m = pf.Model()
>>> m.X = pf.Variable()
>>> m.my_constraint = m.X <= 10
>>> m
<Model vars=1 constrs=1 objective=False>

Try setting the Gurobi license:

>>> m = pf.Model(solver="gurobi", solver_env=dict(ComputeServer="myserver", ServerPassword="mypassword"))
Traceback (most recent call last):
...
RuntimeError: Could not resolve host: myserver (code 6, command POST http://myserver/api/v1/cluster/jobs)

Source code in pyoframe/model.py

def __init__(
    self,
    name=None,
    solver: Optional[SUPPORTED_SOLVER_TYPES] = None,
    solver_env: Optional[Dict[str, str]] = None,
    use_var_names=False,
    sense: Union[ObjSense, ObjSenseValue, None] = None,
):
    self.poi, self.solver_name = Model.create_poi_model(solver, solver_env)
    self._variables: List[Variable] = []
    self._constraints: List[Constraint] = []
    self.sense = ObjSense(sense) if sense is not None else None
    self._objective: Optional[Objective] = None
    self.var_map = (
        NamedVariableMapper(Variable) if Config.print_uses_variable_names else None
    )
    self.name = name

    self.params = Container(self._set_param, self._get_param)
    self.attr = Container(self._set_attr, self._get_attr)
    self._use_var_names = use_var_names

binary_variables property

Examples:

>>> m = pf.Model()
>>> m.X = pf.Variable(vtype=pf.VType.BINARY)
>>> m.Y = pf.Variable()
>>> len(list(m.binary_variables))
1

integer_variables property

Examples:

>>> m = pf.Model()
>>> m.X = pf.Variable(vtype=pf.VType.INTEGER)
>>> m.Y = pf.Variable()
>>> len(list(m.integer_variables))
1

compute_IIS()

Computes the Irreducible Infeasible Set (IIS) of the model.

Gurobi only

This method only works with the Gurobi solver. Open an issue if you'd like to see support for other solvers.

Examples:

>>> m = pf.Model(solver="gurobi")
>>> m.X = pf.Variable(lb=0, ub=2)
>>> m.Y = pf.Variable(lb=0, ub=2)
>>> m.bad_constraint = m.X >= 3
>>> m.minimize = m.X + m.Y
>>> m.optimize()
>>> m.attr.TerminationStatus
<TerminationStatusCode.INFEASIBLE: 3>
>>> m.bad_constraint.attr.IIS
Traceback (most recent call last):
...
polars.exceptions.ComputeError: RuntimeError: Unable to retrieve attribute 'IISConstr'
>>> m.compute_IIS()
>>> m.bad_constraint.attr.IIS
True

Source code in pyoframe/model.py

@for_solvers("gurobi", "copt")
def compute_IIS(self):
    """
    Computes the Irreducible Infeasible Set (IIS) of the model.

    !!! warning "Gurobi only"
        This method only works with the Gurobi solver. Open an issue if you'd like to see support for other solvers.

    Examples:
        >>> m = pf.Model(solver="gurobi")
        >>> m.X = pf.Variable(lb=0, ub=2)
        >>> m.Y = pf.Variable(lb=0, ub=2)
        >>> m.bad_constraint = m.X >= 3
        >>> m.minimize = m.X + m.Y
        >>> m.optimize()
        >>> m.attr.TerminationStatus
        <TerminationStatusCode.INFEASIBLE: 3>
        >>> m.bad_constraint.attr.IIS
        Traceback (most recent call last):
        ...
        polars.exceptions.ComputeError: RuntimeError: Unable to retrieve attribute 'IISConstr'
        >>> m.compute_IIS()
        >>> m.bad_constraint.attr.IIS
        True
    """
    self.poi.computeIIS()

convert_to_fixed()

Turns a mixed integer program into a continuous one by fixing all the integer and binary variables to their solution values.

Gurobi only

This method only works with the Gurobi solver. Open an issue if you'd like to see support for other solvers.

Examples:

>>> m = pf.Model(solver="gurobi")
>>> m.X = pf.Variable(vtype=pf.VType.BINARY, lb=0)
>>> m.Y = pf.Variable(vtype=pf.VType.INTEGER, lb=0)
>>> m.Z = pf.Variable(lb=0)
>>> m.my_constraint = m.X + m.Y + m.Z <= 10
>>> m.maximize = 3 * m.X + 2 * m.Y + m.Z
>>> m.optimize()
>>> m.X.solution, m.Y.solution, m.Z.solution
(1.0, 9.0, 0.0)
>>> m.my_constraint.dual
Traceback (most recent call last):
...
polars.exceptions.ComputeError: RuntimeError: Unable to retrieve attribute 'Pi'
>>> m.convert_to_fixed()
>>> m.optimize()
>>> m.my_constraint.dual
1.0

Only works for Gurobi:

>>> m = pf.Model("max", solver="highs")
>>> m.convert_to_fixed()
Traceback (most recent call last):
...
NotImplementedError: Method 'convert_to_fixed' is not implemented for solver 'highs'.

Source code in pyoframe/model.py

@for_solvers("gurobi")
def convert_to_fixed(self) -> None:
    """
    Turns a mixed integer program into a continuous one by fixing
    all the integer and binary variables to their solution values.

    !!! warning "Gurobi only"
        This method only works with the Gurobi solver. Open an issue if you'd like to see support for other solvers.

    Examples:
        >>> m = pf.Model(solver="gurobi")
        >>> m.X = pf.Variable(vtype=pf.VType.BINARY, lb=0)
        >>> m.Y = pf.Variable(vtype=pf.VType.INTEGER, lb=0)
        >>> m.Z = pf.Variable(lb=0)
        >>> m.my_constraint = m.X + m.Y + m.Z <= 10
        >>> m.maximize = 3 * m.X + 2 * m.Y + m.Z
        >>> m.optimize()
        >>> m.X.solution, m.Y.solution, m.Z.solution
        (1.0, 9.0, 0.0)
        >>> m.my_constraint.dual
        Traceback (most recent call last):
        ...
        polars.exceptions.ComputeError: RuntimeError: Unable to retrieve attribute 'Pi'
        >>> m.convert_to_fixed()
        >>> m.optimize()
        >>> m.my_constraint.dual
        1.0

        Only works for Gurobi:

        >>> m = pf.Model("max", solver="highs")
        >>> m.convert_to_fixed()
        Traceback (most recent call last):
        ...
        NotImplementedError: Method 'convert_to_fixed' is not implemented for solver 'highs'.
    """
    self.poi._converttofixed()

dispose()

Tries to close the solver connection by deleting the model and forcing the garbage collector to run.

Gurobi only. Once this method is called, this model is no longer usable.

This method will not work if you have a variable that references self.poi. Unfortunately, this is a limitation from the underlying solver interface library. See https://github.com/metab0t/PyOptInterface/issues/36 for context.

Examples:

>>> m = pf.Model(solver="gurobi")
>>> m.X = pf.Variable(ub=1)
>>> m.maximize = m.X
>>> m.optimize()
>>> m.X.solution
1.0
>>> m.dispose()
>>> m.X.solution
Traceback (most recent call last):
...
AttributeError: 'Model' object has no attribute 'poi'

Source code in pyoframe/model.py

@for_solvers("gurobi")
def dispose(self):
    """
    Tries to close the solver connection by deleting the model and forcing the garbage collector to run.

    Gurobi only. Once this method is called, this model is no longer usable.

    This method will not work if you have a variable that references self.poi.
    Unfortunately, this is a limitation from the underlying solver interface library.
    See https://github.com/metab0t/PyOptInterface/issues/36 for context.

    Examples:
        >>> m = pf.Model(solver="gurobi")
        >>> m.X = pf.Variable(ub=1)
        >>> m.maximize = m.X
        >>> m.optimize()
        >>> m.X.solution
        1.0
        >>> m.dispose()
        >>> m.X.solution
        Traceback (most recent call last):
        ...
        AttributeError: 'Model' object has no attribute 'poi'
    """
    import gc

    env = self.poi._env
    del self.poi
    gc.collect()
    del env
    gc.collect()

Set(*data, **named_data)

Bases: ModelElement, SupportsMath, SupportPolarsMethodMixin

A set which can then be used to index variables.

Examples:

>>> pf.Set(x=range(2), y=range(3))
<Set size=6 dimensions={'x': 2, 'y': 3}>
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

Source code in pyoframe/core.py

def __init__(self, *data: SetTypes | Iterable[SetTypes], **named_data):
    data_list = list(data)
    for name, set in named_data.items():
        data_list.append({name: set})
    df = self._parse_acceptable_sets(*data_list)
    if not df.is_empty() and df.is_duplicated().any():
        raise ValueError("Duplicate rows found in input data.")
    super().__init__(df)

_parse_acceptable_sets(*over) staticmethod

Examples:

>>> import pandas as pd
>>> dim1 = pd.Index([1, 2, 3], name="dim1")
>>> dim2 = pd.Index(["a", "b"], name="dim1")
>>> Set._parse_acceptable_sets([dim1, dim2])
Traceback (most recent call last):
...
AssertionError: All coordinates must have unique column names.
>>> dim2.name = "dim2"
>>> Set._parse_acceptable_sets([dim1, dim2])
shape: (6, 2)
┌──────┬──────┐
│ dim1 ┆ dim2 │
│ ---  ┆ ---  │
│ i64  ┆ str  │
╞══════╪══════╡
│ 1    ┆ a    │
│ 1    ┆ b    │
│ 2    ┆ a    │
│ 2    ┆ b    │
│ 3    ┆ a    │
│ 3    ┆ b    │
└──────┴──────┘

Source code in pyoframe/core.py

@staticmethod
def _parse_acceptable_sets(
    *over: SetTypes | Iterable[SetTypes],
) -> pl.DataFrame:
    """
    Examples:
        >>> import pandas as pd
        >>> dim1 = pd.Index([1, 2, 3], name="dim1")
        >>> dim2 = pd.Index(["a", "b"], name="dim1")
        >>> Set._parse_acceptable_sets([dim1, dim2])
        Traceback (most recent call last):
        ...
        AssertionError: All coordinates must have unique column names.
        >>> dim2.name = "dim2"
        >>> Set._parse_acceptable_sets([dim1, dim2])
        shape: (6, 2)
        ┌──────┬──────┐
        │ dim1 ┆ dim2 │
        │ ---  ┆ ---  │
        │ i64  ┆ str  │
        ╞══════╪══════╡
        │ 1    ┆ a    │
        │ 1    ┆ b    │
        │ 2    ┆ a    │
        │ 2    ┆ b    │
        │ 3    ┆ a    │
        │ 3    ┆ b    │
        └──────┴──────┘
    """
    assert len(over) > 0, "At least one set must be provided."
    over_iter: Iterable[SetTypes] = parse_inputs_as_iterable(*over)

    over_frames: List[pl.DataFrame] = [Set._set_to_polars(set) for set in over_iter]

    over_merged = over_frames[0]

    for df in over_frames[1:]:
        assert (
            set(over_merged.columns) & set(df.columns) == set()
        ), "All coordinates must have unique column names."
        over_merged = over_merged.join(df, how="cross")
    return over_merged

Variable(*indexing_sets, lb=None, ub=None, vtype=VType.CONTINUOUS, equals=None)

Bases: ModelElementWithId, SupportsMath, SupportPolarsMethodMixin

Represents one or many decision variable in an optimization model.

Parameters:

Name	Type	Description	Default
*indexing_sets	SetTypes | Iterable[SetTypes]	If no indexing_sets are provided, a single variable with no dimensions is created. Otherwise, a variable is created for each element in the Cartesian product of the indexing_sets (see Set for details on behaviour).	()
lb	float | int | SupportsToExpr | None	The lower bound for all variables.	None
ub	float | int | SupportsToExpr | None	The upper bound for all variables.	None
vtype	VType | VTypeValue	The type of the variable. Can be either a VType enum or a string. Default is VType.CONTINUOUS.	CONTINUOUS
equals	Optional[SupportsMath]	When specified, a variable is created and a constraint is added to make the variable equal to the provided expression.	None

Examples:

>>> import pandas as pd
>>> m = pf.Model()
>>> df = pd.DataFrame({"dim1": [1, 1, 2, 2, 3, 3], "dim2": ["a", "b", "a", "b", "a", "b"]})
>>> v = Variable(df)
>>> v
<Variable size=6 dimensions={'dim1': 3, 'dim2': 2} added_to_model=False>

Variables cannot be used until they're added to the model.

>>> m.constraint = v <= 3
Traceback (most recent call last):
...
ValueError: Cannot use 'Variable' before it has beed added to a model.
>>> m.v = v
>>> m.constraint = m.v <= 3

>>> m.v
<Variable name=v size=6 dimensions={'dim1': 3, 'dim2': 2}>
[1,a]: v[1,a]
[1,b]: v[1,b]
[2,a]: v[2,a]
[2,b]: v[2,b]
[3,a]: v[3,a]
[3,b]: v[3,b]
>>> m.v2 = Variable(df[["dim1"]])
Traceback (most recent call last):
...
ValueError: Duplicate rows found in input data.
>>> m.v3 = Variable(df[["dim1"]].drop_duplicates())
>>> m.v3
<Variable name=v3 size=3 dimensions={'dim1': 3}>
[1]: v3[1]
[2]: v3[2]
[3]: v3[3]

Source code in pyoframe/core.py

def __init__(
    self,
    *indexing_sets: SetTypes | Iterable[SetTypes],
    lb: float | int | SupportsToExpr | None = None,
    ub: float | int | SupportsToExpr | None = None,
    vtype: VType | VTypeValue = VType.CONTINUOUS,
    equals: Optional[SupportsMath] = None,
):
    if equals is not None:
        assert (
            len(indexing_sets) == 0
        ), "Cannot specify both 'equals' and 'indexing_sets'"
        indexing_sets = (equals,)

    data = Set(*indexing_sets).data if len(indexing_sets) > 0 else pl.DataFrame()
    super().__init__(data)

    self.vtype: VType = VType(vtype)
    self.attr = Container(self._set_attribute, self._get_attribute)
    self._equals = equals

    if lb is not None and not isinstance(lb, (float, int)):
        self._lb_expr, self.lb = lb, None
    else:
        self._lb_expr, self.lb = None, lb
    if ub is not None and not isinstance(ub, (float, int)):
        self._ub_expr, self.ub = ub, None
    else:
        self._ub_expr, self.ub = None, ub

next(dim, wrap_around=False)

Creates an expression where the variable at each index is the next variable in the specified dimension.

Parameters:

Name	Type	Description	Default
dim	str	The dimension over which to shift the variable.	required
wrap_around	bool	If True, the last index in the dimension is connected to the first index.	False

Examples:

>>> import pandas as pd
>>> time_dim = pd.DataFrame({"time": ["00:00", "06:00", "12:00", "18:00"]})
>>> space_dim = pd.DataFrame({"city": ["Toronto", "Berlin"]})
>>> m = pf.Model()
>>> m.bat_charge = pf.Variable(time_dim, space_dim)
>>> m.bat_flow = pf.Variable(time_dim, space_dim)
>>> # Fails because the dimensions are not the same
>>> m.bat_charge + m.bat_flow == m.bat_charge.next("time")
Traceback (most recent call last):
...
pyoframe.constants.PyoframeError: Failed to add expressions:
<Expression size=8 dimensions={'time': 4, 'city': 2} terms=16> + <Expression size=6 dimensions={'city': 2, 'time': 3} terms=6>
Due to error:
Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()
shape: (2, 4)
┌───────┬─────────┬────────────┬────────────┐
│ time  ┆ city    ┆ time_right ┆ city_right │
│ ---   ┆ ---     ┆ ---        ┆ ---        │
│ str   ┆ str     ┆ str        ┆ str        │
╞═══════╪═════════╪════════════╪════════════╡
│ 18:00 ┆ Toronto ┆ null       ┆ null       │
│ 18:00 ┆ Berlin  ┆ null       ┆ null       │
└───────┴─────────┴────────────┴────────────┘

>>> (m.bat_charge + m.bat_flow).drop_unmatched() == m.bat_charge.next("time")
<Constraint sense='=' size=6 dimensions={'time': 3, 'city': 2} terms=18>
[00:00,Berlin]: bat_charge[00:00,Berlin] + bat_flow[00:00,Berlin] - bat_charge[06:00,Berlin] = 0
[00:00,Toronto]: bat_charge[00:00,Toronto] + bat_flow[00:00,Toronto] - bat_charge[06:00,Toronto] = 0
[06:00,Berlin]: bat_charge[06:00,Berlin] + bat_flow[06:00,Berlin] - bat_charge[12:00,Berlin] = 0
[06:00,Toronto]: bat_charge[06:00,Toronto] + bat_flow[06:00,Toronto] - bat_charge[12:00,Toronto] = 0
[12:00,Berlin]: bat_charge[12:00,Berlin] + bat_flow[12:00,Berlin] - bat_charge[18:00,Berlin] = 0
[12:00,Toronto]: bat_charge[12:00,Toronto] + bat_flow[12:00,Toronto] - bat_charge[18:00,Toronto] = 0

>>> (m.bat_charge + m.bat_flow) == m.bat_charge.next("time", wrap_around=True)
<Constraint sense='=' size=8 dimensions={'time': 4, 'city': 2} terms=24>
[00:00,Berlin]: bat_charge[00:00,Berlin] + bat_flow[00:00,Berlin] - bat_charge[06:00,Berlin] = 0
[00:00,Toronto]: bat_charge[00:00,Toronto] + bat_flow[00:00,Toronto] - bat_charge[06:00,Toronto] = 0
[06:00,Berlin]: bat_charge[06:00,Berlin] + bat_flow[06:00,Berlin] - bat_charge[12:00,Berlin] = 0
[06:00,Toronto]: bat_charge[06:00,Toronto] + bat_flow[06:00,Toronto] - bat_charge[12:00,Toronto] = 0
[12:00,Berlin]: bat_charge[12:00,Berlin] + bat_flow[12:00,Berlin] - bat_charge[18:00,Berlin] = 0
[12:00,Toronto]: bat_charge[12:00,Toronto] + bat_flow[12:00,Toronto] - bat_charge[18:00,Toronto] = 0
[18:00,Berlin]: bat_charge[18:00,Berlin] + bat_flow[18:00,Berlin] - bat_charge[00:00,Berlin] = 0
[18:00,Toronto]: bat_charge[18:00,Toronto] + bat_flow[18:00,Toronto] - bat_charge[00:00,Toronto] = 0

Source code in pyoframe/core.py

def next(self, dim: str, wrap_around: bool = False) -> Expression:
    """
    Creates an expression where the variable at each index is the next variable in the specified dimension.

    Parameters:
        dim:
            The dimension over which to shift the variable.
        wrap_around:
            If True, the last index in the dimension is connected to the first index.

    Examples:
        >>> import pandas as pd
        >>> time_dim = pd.DataFrame({"time": ["00:00", "06:00", "12:00", "18:00"]})
        >>> space_dim = pd.DataFrame({"city": ["Toronto", "Berlin"]})
        >>> m = pf.Model()
        >>> m.bat_charge = pf.Variable(time_dim, space_dim)
        >>> m.bat_flow = pf.Variable(time_dim, space_dim)
        >>> # Fails because the dimensions are not the same
        >>> m.bat_charge + m.bat_flow == m.bat_charge.next("time")
        Traceback (most recent call last):
        ...
        pyoframe.constants.PyoframeError: Failed to add expressions:
        <Expression size=8 dimensions={'time': 4, 'city': 2} terms=16> + <Expression size=6 dimensions={'city': 2, 'time': 3} terms=6>
        Due to error:
        Dataframe has unmatched values. If this is intentional, use .drop_unmatched() or .keep_unmatched()
        shape: (2, 4)
        ┌───────┬─────────┬────────────┬────────────┐
        │ time  ┆ city    ┆ time_right ┆ city_right │
        │ ---   ┆ ---     ┆ ---        ┆ ---        │
        │ str   ┆ str     ┆ str        ┆ str        │
        ╞═══════╪═════════╪════════════╪════════════╡
        │ 18:00 ┆ Toronto ┆ null       ┆ null       │
        │ 18:00 ┆ Berlin  ┆ null       ┆ null       │
        └───────┴─────────┴────────────┴────────────┘

        >>> (m.bat_charge + m.bat_flow).drop_unmatched() == m.bat_charge.next("time")
        <Constraint sense='=' size=6 dimensions={'time': 3, 'city': 2} terms=18>
        [00:00,Berlin]: bat_charge[00:00,Berlin] + bat_flow[00:00,Berlin] - bat_charge[06:00,Berlin] = 0
        [00:00,Toronto]: bat_charge[00:00,Toronto] + bat_flow[00:00,Toronto] - bat_charge[06:00,Toronto] = 0
        [06:00,Berlin]: bat_charge[06:00,Berlin] + bat_flow[06:00,Berlin] - bat_charge[12:00,Berlin] = 0
        [06:00,Toronto]: bat_charge[06:00,Toronto] + bat_flow[06:00,Toronto] - bat_charge[12:00,Toronto] = 0
        [12:00,Berlin]: bat_charge[12:00,Berlin] + bat_flow[12:00,Berlin] - bat_charge[18:00,Berlin] = 0
        [12:00,Toronto]: bat_charge[12:00,Toronto] + bat_flow[12:00,Toronto] - bat_charge[18:00,Toronto] = 0

        >>> (m.bat_charge + m.bat_flow) == m.bat_charge.next("time", wrap_around=True)
        <Constraint sense='=' size=8 dimensions={'time': 4, 'city': 2} terms=24>
        [00:00,Berlin]: bat_charge[00:00,Berlin] + bat_flow[00:00,Berlin] - bat_charge[06:00,Berlin] = 0
        [00:00,Toronto]: bat_charge[00:00,Toronto] + bat_flow[00:00,Toronto] - bat_charge[06:00,Toronto] = 0
        [06:00,Berlin]: bat_charge[06:00,Berlin] + bat_flow[06:00,Berlin] - bat_charge[12:00,Berlin] = 0
        [06:00,Toronto]: bat_charge[06:00,Toronto] + bat_flow[06:00,Toronto] - bat_charge[12:00,Toronto] = 0
        [12:00,Berlin]: bat_charge[12:00,Berlin] + bat_flow[12:00,Berlin] - bat_charge[18:00,Berlin] = 0
        [12:00,Toronto]: bat_charge[12:00,Toronto] + bat_flow[12:00,Toronto] - bat_charge[18:00,Toronto] = 0
        [18:00,Berlin]: bat_charge[18:00,Berlin] + bat_flow[18:00,Berlin] - bat_charge[00:00,Berlin] = 0
        [18:00,Toronto]: bat_charge[18:00,Toronto] + bat_flow[18:00,Toronto] - bat_charge[00:00,Toronto] = 0
    """

    wrapped = self.data.select(dim).unique(maintain_order=True).sort(by=dim)
    wrapped = wrapped.with_columns(pl.col(dim).shift(-1).alias("__next"))
    if wrap_around:
        wrapped = wrapped.with_columns(pl.col("__next").fill_null(pl.first(dim)))
    else:
        wrapped = wrapped.drop_nulls(dim)

    expr = self.to_expr()
    data = expr.data.rename({dim: "__prev"})

    if POLARS_VERSION.major < 1:
        data = data.join(
            wrapped, left_on="__prev", right_on="__next", how="inner"
        ).drop(["__prev", "__next"])
    else:
        data = data.join(
            wrapped, left_on="__prev", right_on="__next", how="inner"
        ).drop(["__prev", "__next"], strict=False)
    return expr._new(data)