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). Phew, all in one sentence! Hope you made it til here.

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.

Edit: Result

The answer and comment triggered me to explore how much the constraint definition can be split. I ended up writing a helper function:

def commodity_balance(m, tm, co):
    """ calculate commodity balance at given timestep.
    
    [more docstring]"""
    balance = 0
    for p in m.pro_tuples:
        if p[1] == co:
            # usage as input for process increases balance
            balance += m.e_pro_in[(tm,)+p]
        if p[2] == co:
            # output from processes decreases balance
            balance -= m.e_pro_out[(tm,)+p]
    for s in m.sto_tuples:
        # usage as input for storage increases consumption
        # output from storage decreases consumption
        if s[1] == co:
            balance += m.e_sto_in[(tm,)+s]
            balance -= m.e_sto_out[(tm,)+s]
    return balance

With its help, the ugly res_stock_total_rule becomes a breeze:

def res_stock_total_rule(m, co, co_type):
    if co not in m.co_stock:
        return Constraint.Skip
    else:
        # calculate total consumption of commodity co
        total_consumption = 0
        for tm in m.tm:
            total_consumption += commodity_balance(m, tm, co)                
         
        return total_consumption <= m.commodity.loc[co, co_type]['max']

Bonus: I can reuse the function for two other rules, which happened to be the only constraints for which I had readability concerns. Thanks!