Write constraints with relationship sets in Pyomo

Question

I would like to write the following constraint in the most compact form possible using Pyomo.

$$\sum_{o\in O}y_{n,o} = \sum_{o\in O} \sum_{n'\in LN}y_{n',o}\qquad \forall n$$

Where, the important part, is that $LN$ is a relationship set, or multiset, such that its elements are tuples like $(n_1,n_2), (n_2, n_4), ...$

In GAMS, this is easy. The code for a small example would look like:

SETS
n       /n1*n6/
LN(n,n) /(n1.n2), (n2.n4)/
O       /1,2,3/
;
alias(n,nn);

VARIABLES
y(n,o)
;

EQUATIONS
EQ01;
EQ01(n)..      sum(O, y(n,o)) = sum(o, sum(nn$LN(n,nn), y(nn,o));

For example, for $n=n1$, I would want the following equation. $$y_{n_1,1} + y_{n_1,2} + y_{n_1,3} = y_{n_2,1} + y_{n_2,2} + y_{n_2,3}$$

I'm not sure of how to translate the compactness of GAMS to Pyomo (if possible). I am working with some rules for a pyomo model in the form:

model    = pyo.ConcreteModel()
model.N  = pyo.Set(initialize = ["n1","n2","n3","n4","n5", "n6"]
model.LN = pyo.Set(within = model.N*model.N, initialize = [("n1", "n2"), ("n2", "n4")]
model.O  = pyo.Set(initialize = [1,2,3])
model.y  = pyo.Var(model.N, model.O)

def _rule(model,n):
  return sum(model.y[n,o] for o in model.O) == sum(sum(model.y[nn,o] for nn in model.LN) for o in model.O) 
model.equation = pyo.Constraint(model.N, rule = _rule)

But it is clearly not working, since in the inner sum it does not know what nn is. The set of LN also doesn't have a way of knowing that I'm referring to the outer $n$ in there, I think.

I think something could be done with the advantage of Pyomo's ability to use "ifs" and other flow controls, together with ConstraintSkip for the cases where there is no $n'$ associated to a $n$. However, I would like to know if there is a more elegant, compact way of writing this, such as GAMS'.

Thank you for your time!

Answer 1

First create a dictionary like

Allowed={"n1":"n2", "n2":"n4"}

and then

Use this

from pyomo.environ import *
Allowed={"n1":"n2", "n2":"n4"}
model    = AbstractModel()
model.N  = Set(initialize = ["n1","n2","n3","n4"])
model.O  = Set(initialize = [1,2,3])
model.y  = Var(model.N, model.O, within=Reals)
def yrule(model,n):
    if n in Allowed:
        return sum(model.y[n,o] for o in model.O) ==sum(model.y[nn,o] for nn in model.N for o in model.O if nn in Allowed[n])
    else:
        return Constraint.Skip
model.equation = Constraint(model.N, rule = yrule)

instance = model.create_instance()
instance.equation.pprint()

Answer 2

I will mark the answer from Optimization team, which is completely correct, as it works for the problem given, although it does not use tuple sets. I will also put here for completion how I ended up doing it, since I found that it could be made more compact with list comprehension and without Constraint.Skip.

Apparently, I was on the right track with this:

def _rule(model,n):
  return sum(model.y[n,o] for o in model.O) == sum(sum(model.y[nn,o] for nn in model.LN) for o in model.O)

But I had to slightly change it to this.

def _rule(model,n):
  return sum(model.y[n,o] for o in model.O) == sum(sum(model.y[nn,o] for nn in m.N if (n,nn) in model.LN) for o in model.O)

So now it knows where the nn should come from, the set N, and also knows that only do it for the pairs that are present in the tuple set.

Thank you all who helped me!