Skip to content

core

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))

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

SupportsMath(**kwargs)

Bases: ABC, SupportsToExpr

Any object that can be converted into an expression.

Source code in pyoframe/core.py
def __init__(self, **kwargs):
    self.unmatched_strategy = UnmatchedStrategy.UNSET
    self.allowed_new_dims: List[str] = []
    super().__init__(**kwargs)

__eq__(value)

Equality constraint.

Examples:

>>> m = pf.Model()
>>> m.v = pf.Variable()
>>> m.v == 1
<Constraint sense='=' size=1 dimensions={} terms=2>
v = 1
Source code in pyoframe/core.py
def __eq__(self, value: object):  # type: ignore
    """Equality constraint.

    Examples:
        >>> m = pf.Model()
        >>> m.v = pf.Variable()
        >>> m.v == 1
        <Constraint sense='=' size=1 dimensions={} terms=2>
        v = 1
    """
    return Constraint(self - value, ConstraintSense.EQ)

__ge__(other)

Equality constraint.

Examples:

>>> m = pf.Model()
>>> m.v = pf.Variable()
>>> m.v >= 1
<Constraint sense='>=' size=1 dimensions={} terms=2>
v >= 1
Source code in pyoframe/core.py
def __ge__(self, other):
    """Equality constraint.

    Examples:
        >>> m = pf.Model()
        >>> m.v = pf.Variable()
        >>> m.v >= 1
        <Constraint sense='>=' size=1 dimensions={} terms=2>
        v >= 1
    """
    return Constraint(self - other, ConstraintSense.GE)

__le__(other)

Equality constraint.

Examples:

>>> m = pf.Model()
>>> m.v = pf.Variable()
>>> m.v <= 1
<Constraint sense='<=' size=1 dimensions={} terms=2>
v <= 1
Source code in pyoframe/core.py
def __le__(self, other):
    """Equality constraint.

    Examples:
        >>> m = pf.Model()
        >>> m.v = pf.Variable()
        >>> m.v <= 1
        <Constraint sense='<=' size=1 dimensions={} terms=2>
        v <= 1
    """
    return Constraint(self - other, ConstraintSense.LE)

__pow__(power)

Support squaring expressions:

m = pf.Model() m.v = pf.Variable() m.v ** 2 v * v m.v ** 3 Traceback (most recent call last): ... ValueError: Raising an expressions to 3 is not supported. Expressions can only be squared (2).

Source code in pyoframe/core.py
def __pow__(self, power: int):
    """
    Support squaring expressions:
    >>> m = pf.Model()
    >>> m.v = pf.Variable()
    >>> m.v ** 2
    <Expression size=1 dimensions={} terms=1 degree=2>
    v * v
    >>> m.v ** 3
    Traceback (most recent call last):
    ...
    ValueError: Raising an expressions to **3 is not supported. Expressions can only be squared (**2).
    """
    if power == 2:
        return self * self
    raise ValueError(
        f"Raising an expressions to **{power} is not supported. Expressions can only be squared (**2)."
    )

__rsub__(other)

Support right subtraction.

Examples:

>>> m = pf.Model()
>>> m.v = Variable({"dim1": [1,2,3]})
>>> 1 - m.v
<Expression size=3 dimensions={'dim1': 3} terms=6>
[1]: 1  - v[1]
[2]: 1  - v[2]
[3]: 1  - v[3]
Source code in pyoframe/core.py
def __rsub__(self, other):
    """
    Support right subtraction.

    Examples:
        >>> m = pf.Model()
        >>> m.v = Variable({"dim1": [1,2,3]})
        >>> 1 - m.v
        <Expression size=3 dimensions={'dim1': 3} terms=6>
        [1]: 1  - v[1]
        [2]: 1  - v[2]
        [3]: 1  - v[3]
    """
    return other + (-self.to_expr())

__sub__(other)

import polars as pl m = pf.Model() df = pl.DataFrame({"dim1": [1,2,3], "value": [1,2,3]}) m.v = pf.Variable(df["dim1"]) m.v - df [1]: v[1] -1 [2]: v[2] -2 [3]: v[3] -3

Source code in pyoframe/core.py
def __sub__(self, other):
    """
    >>> import polars as pl
    >>> m = pf.Model()
    >>> df = pl.DataFrame({"dim1": [1,2,3], "value": [1,2,3]})
    >>> m.v = pf.Variable(df["dim1"])
    >>> m.v - df
    <Expression size=3 dimensions={'dim1': 3} terms=6>
    [1]: v[1] -1
    [2]: v[2] -2
    [3]: v[3] -3
    """
    if not isinstance(other, (int, float)):
        other = other.to_expr()
    return self.to_expr() + (-other)

__truediv__(other)

Examples:

Support division.

>>> m = pf.Model()
>>> m.v = Variable({"dim1": [1,2,3]})
>>> m.v / 2
<Expression size=3 dimensions={'dim1': 3} terms=3>
[1]: 0.5 v[1]
[2]: 0.5 v[2]
[3]: 0.5 v[3]
Source code in pyoframe/core.py
def __truediv__(self, other):
    """

    Examples:
        Support division.
        >>> m = pf.Model()
        >>> m.v = Variable({"dim1": [1,2,3]})
        >>> m.v / 2
        <Expression size=3 dimensions={'dim1': 3} terms=3>
        [1]: 0.5 v[1]
        [2]: 0.5 v[2]
        [3]: 0.5 v[3]
    """
    return self.to_expr() * (1 / other)

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)