from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING
import gamspy as gp
from gamspy._algebra.condition import Condition
from gamspy._algebra.domain import Domain
from gamspy._symbols.implicits import ImplicitSet
from gamspy.exceptions import ValidationError
if TYPE_CHECKING:
from gamspy import Alias, Container, Set
from gamspy._algebra.expression import Expression
[docs]
class Loop:
"""
A context manager to execute a group of statements iteratively for each member of a set or domain.
The Loop class maps to the GAMS `loop` statement. It is particularly useful for
cases where parallel assignments are not sufficient, such as iterative calculations,
nested loops, or modifying models and solving them repeatedly.
Parameters
----------
indices : Set | Alias | ImplicitSet | Condition | Domain | Sequence[Set | Alias]
The controlling domain of the loop. This can be a single Set, a sequence of Sets,
or a domain restricted by a logical condition (using `.where`).
Examples
--------
**1. Simple iteration over a single Set:**
>>> import gamspy as gp
>>> m = gp.Container()
>>> t = gp.Set(m, records=["1985", "1986", "1987"])
>>> pop = gp.Parameter(m, domain=t, records=[("1985", 3456)])
>>> growth = gp.Parameter(m, domain=t, records=[("1985", 25.3), ("1986", 27.3)])
>>> with gp.Loop(t):
... pop[t + 1] = pop[t] + growth[t]
**2. Iteration with a logical condition (dollar condition):**
You can restrict the loop domain using the `.where` attribute on Sets or Domains.
>>> i = gp.Set(m, records=["i1", "i2", "i3"])
>>> j = gp.Set(m, records=["j1", "j2", "j3"])
>>> q = gp.Parameter(m, domain=[i, j], records=[("i1", "j1", 1), ("i1", "j2", 3)])
>>> x = gp.Parameter(m, records=1)
>>> with gp.Loop(gp.Domain(i, j).where[q[i, j] > 0]):
... x[...] = x[...] + q[i, j]
**3. Nested Loops:**
Loops can be nested using standard Python indentation.
>>> a = gp.Parameter(m, domain=[i, j])
>>> b = gp.Parameter(m)
>>> a.generateRecords()
>>> with gp.Loop(i):
... with gp.Loop(j):
... b[...] = a[i, j]
"""
def __init__(
self,
indices: Set | Alias | ImplicitSet | Condition | Domain | Sequence[Set | Alias],
):
self.indices = indices
self._loop_number = -1
self.container = self._find_container()
def _find_container(self) -> Container:
if isinstance(self.indices, (gp.Set, gp.Alias, Condition, Domain, ImplicitSet)):
return self.indices.container # type: ignore
elif isinstance(self.indices, Sequence):
for elem in self.indices:
if hasattr(elem, "container"):
return elem.container
raise ValidationError(
f"`{type(self.indices)}` is not an allowed type for a loop index. "
)
def _index_repr(self) -> str:
if isinstance(self.indices, (gp.Set, gp.Alias, Condition, Domain, ImplicitSet)):
return self.indices.gamsRepr()
elif isinstance(self.indices, Sequence):
representations = [index.gamsRepr() for index in self.indices]
return f"({','.join(representations)})"
raise ValidationError(
f"`{type(self.indices)}` is not an allowed type for a loop index. "
)
@property
def Break(self) -> None:
"""
Breaks the execution of the current loop prematurely.
This property maps to the GAMS `break` statement. Note that you can only
break out of the innermost loop currently executing. Attempting to break an
outer loop from within an inner loop will raise a ValidationError.
Raises
------
ValidationError
If attempting to break an outer loop without breaking the inner loop first.
Examples
--------
>>> import gamspy as gp
>>> m = gp.Container()
>>> i = gp.Set(m, records=["i1", "i2", "i3"])
>>> j = gp.Set(m, records=["j1", "j2", "j3"])
>>> cnt = gp.Parameter(m, records=0)
>>> with gp.Loop(i) as loop:
... with gp.Loop(j) as loop2:
... cnt[...] += 1
... loop2.Break # Successfully breaks the inner loop
... loop.Break # Successfully breaks the outer loop
"""
if self._loop_number < self.container._in_loop:
raise ValidationError(
"You cannot break this loop. You should break the inner loop first."
)
self.container._add_statement("break;")
@property
def Continue(self) -> None:
"""
Skips the remaining statements in the current iteration and proceeds to the next one.
This property maps to the GAMS `continue` statement. It gives additional
control over the execution of loop structures by allowing you to bypass
the rest of the loop block for the current domain element.
Examples
--------
>>> import gamspy as gp
>>> m = gp.Container()
>>> i = gp.Set(m, records=["i1", "i2", "i3", "i4"])
>>> cnt = gp.Parameter(m, records=0)
>>> with gp.Loop(i) as loop:
... with gp.If(gp.Ord(i) == 2):
... loop.Continue # Skips incrementing for "i2"
... cnt[...] += 1
"""
self.container._add_statement("continue;")
def __enter__(self) -> Loop:
self.container._in_loop += 1
self._loop_number = self.container._in_loop
self.container._add_statement(f"loop({self._index_repr()},")
return self
def __exit__(self, exc_type, exc, tb):
self.container._in_loop -= 1
self.container._add_statement(");")
if self.container._in_loop == 0: # Run only in the most outer loop
self.container._synch_with_gams(gams_to_gamspy=True)
[docs]
class If:
"""
A context manager to conditionally execute a group of statements.
The If class maps to the GAMS `if` statement. It allows you to branch
conditionally around a group of execution statements within a loop.
Parameters
----------
condition : Expression
The logical condition that must be satisfied to execute the nested statements.
Examples
--------
**1. Skipping iterations conditionally:**
>>> import gamspy as gp
>>> m = gp.Container()
>>> i = gp.Set(m, records=[f"i{idx}" for idx in range(1, 11)])
>>> cnt = gp.Parameter(m, records=0)
>>> with gp.Loop(i) as loop:
... with gp.If(gp.math.mod(gp.Ord(i), 2) == 0):
... loop.Continue
... cnt[...] += 1
**2. Breaking a loop based on a condition:**
>>> with gp.Loop(i) as loop:
... with gp.If(i.sameAs("i6")):
... loop.Break
... cnt[...] += 1
"""
def __init__(self, condition: Expression | Condition):
self.condition = condition
self.container = condition.container
def __enter__(self):
if not self.container._in_loop:
raise ValidationError(
"`gp.If` context manager can only be used in `gp.Loop` context managers. Use regular Python if statements instead."
)
representation = self.condition.gamsRepr()
representation = gp.utils._replace_equality_signs(representation)
self.container._add_statement(f"if ({representation},")
def __exit__(self, exc_type, exc, tb):
self.container._add_statement(");")
