[Tutor] Question in regards to loops and matplotlib

[Peter Otten]

Zachary Rizer wrote:

> def make_mass_cuts(table):
>     """
>     make mass cuts using logMass row and creates four-tuple
>     each index is divided into its own dual tuple, x and y UVJ colors
>     array returned = [(m1_x, m1_y), (m2_x, m2_y), (m3_x, m3_y), (m4_x,
> m4_y)]
>     """
>     m1_x = []
>     m1_y = []
>     m2_x = []
>     m2_y = []
>     m3_x = []
>     m3_y = []
>     m4_x = []
>     m4_y = []
>     for row in table:
>         if row['ms_lmass'] >= 9.7 and row['ms_lmass'] < 10:
>             m1_x.append(row['rf_VminJ'])
>             m1_y.append(row['rf_UminV'])
>         elif row['ms_lmass'] >= 10 and row['ms_lmass'] < 10.5:
>             m2_x.append(row['rf_VminJ'])
>             m2_y.append(row['rf_UminV'])
>         elif row['ms_lmass'] >= 10.5 and row['ms_lmass'] < 10.8:
>             m3_x.append(row['rf_VminJ'])
>             m3_y.append(row['rf_UminV'])
>         elif row['ms_lmass'] >= 10.8:
>             m4_x.append(row['rf_VminJ'])
>             m4_y.append(row['rf_UminV'])
>     return [(m1_x, m1_y), (m2_x, m2_y), (m3_x, m3_y), (m4_x, m4_y)]

You have redundant tests there:

>         if row['ms_lmass'] >= 9.7 and row['ms_lmass'] < 10:
             ...
>         elif row['ms_lmass'] >= 10 and row['ms_lmass'] < 10.5:

At this point you already know that the mass is >= 10, no need to repeat the 
check. Which gives:

def _make_mass_cuts(table):
    bins = [([], []) for i in range(4)]
    for row in table:
        mass = row["ms_lmass"]
        if mass < 9.7:
            continue
        elif mass < 10:
            bin_index =  0
        elif mass < 10.5:
            bin_index = 1
        elif mass < 10.8:
            bin_index = 2
        else:
            bin_index = 3
        binx, biny = bins[bin_index]
        binx.append(row['rf_VminJ'])
        biny.append(row['rf_UminV'])

    return bins

For so few cuts it's probably not worthwhile, but the bisect module allows 
to generalise this:

import bisect
def make_mass_cuts(table):
    cuts = [9.7, 10.0, 10.5, 10.8]
    bins = [([], []) for cut in cuts]
    for row in table:
        index = bisect.bisect_right(cuts, row['ms_lmass'])
        if index:
            x, y = bins[index-1]
            x.append(row['rf_VminJ'])
            y.append(row['rf_UminV'])
    return bins

If you want to play with the values or the number of bins now all you have 
to change is the cuts list.