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.

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!

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

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!