Long expression to sum data in a multi-dimensional model

Question

I am porting a linear optimization model for power plants from GAMS to Pyomo. Models in both frameworks are a collection of sets (both elementary or tuple sets), parameters (fixed values, defined over sets), variables (unknowns, defined over sets, value to be determined by optimization) and equations (defining relationships between variables and parameters).

In the following example, I am asking for ideas on how to make the following inequality more readable:

def res_stock_total_rule(m, co, co_type):
    if co in m.co_stock:
        return sum(m.e_pro_in[(tm,)+ p] for tm in m.tm for p in m.pro_tuples if p[1] == co) + \
               sum(m.e_pro_out[(tm,)+ p] for tm in m.tm for p in m.pro_tuples if p[2] == co) + \
               sum(m.e_sto_out[(tm,)+ s] for tm in m.tm for s in m.sto_tuples if s[1] == co) - \
               sum(m.e_sto_in[(tm,)+ s] for tm in m.tm for s in m.sto_tuples if s[1] == co) <= \
               m.commodity.loc[co, co_type]['max']
    else:
        return Constraint.Skip

Context:

• m is a model object, which contains all of the above elements (sets, params, variables, equations) as attributes.
• m.e_pro_in for example is a 4-dimensional variable defined over the tuple set (time, process name, input commodity, output commodity).
• m.tm is a set of timesteps t = {1, 2, ...}, m.co_stock the set of stock commodity, for which this rule will apply only (otherwise, no Constraint is generated via Skip).
• m.pro_tuples is a set of all valid (i.e. realisable) tuples (process name, input commodity, output commodity).
• m.commodity is a Pandas DataFrame that effectively acts as a model parameter.

My question now is this:

Can you give me some hints on how to improve the readability of this fragment? The combination of tuple concatenation, two nested list comprehensions with conditional clause, Pandas DataFrame indexing, and a multiline expression with line breaks all make it less than easy to read for someone who might just be learning Python while using this model.

Answer 1

First of all, use helper functions or explicit for loops.

E.g. you're looping over m.tm four times. This can be done in a single loop. It might need more lines of code, but it will get much more readable.

E.g.

def res_stock_total_rule(m, co, co_type):
    if not co in m.co_stock:
        return Constraint.Skip

    val = 0
    for tm in m.tm: # one single loop instead of four
        for p in m.pro_tuples:
            if p[1] == co:
                val += m.e_pro_in[(tm,) + p)]
            if p[2] == co:
                val += m.e_pro_out[(tm,) + p)]
        for s in m.sto_tuples: # one single loop instead of two
            if s[1] == co:
                val += m.e_sto_out[(tm,) + p)] - m.e_sto_in[(tm,) + p)]
    return val <= m.commodity.loc[co, co_type]['max']

Answer 2

Wrapping with parentheses is often preferable to newline escapes:

def res_stock_total_rule(m, co, co_type):
    if co in m.co_stock:
        return (
            sum(m.e_pro_in [(tm,) + p] for tm in m.tm for p in m.pro_tuples if p[1] == co) +
            sum(m.e_pro_out[(tm,) + p] for tm in m.tm for p in m.pro_tuples if p[2] == co) +
            sum(m.e_sto_out[(tm,) + s] for tm in m.tm for s in m.sto_tuples if s[1] == co) -
            sum(m.e_sto_in [(tm,) + s] for tm in m.tm for s in m.sto_tuples if s[1] == co)
        ) <= m.commodity.loc[co, co_type]['max']
    else:
        return Constraint.Skip