Python - masking in a for loop?

Question

I have three arrays, r_vals, Tgas_vals, and n_vals. They are all numpy arrays of the shape (9998.). The arrays have repeated values and I want to iterate over the unique values of r_vals and find the corresponding values of Tgas_vals, and n_vals so I can use the last two arrays to calculate the weighted average. This is what I have right now:

def calc_weighted_average (r_vals,Tgas_vals,n_vals):
        for r in r_vals:
            mask = r == r_vals
            count = 0
            count += 1
            for t in Tgas_vals[mask]:
                print (count, np.average(Tgas_vals[mask]*n_vals[mask]))

    weighted_average = calc_weighted_average (r_vals,Tgas_vals,n_vals)

The problem I am running into is that the function is only looping through once. Did I implement mask incorrectly, or is the problem somewhere else in the for loop?

You'll have elaborate your question. Is `unique_r` actually `unique(r_vals)`. I don't see where `z_vals` is needed at all in your code either. Also, you should fix your indentation too. — slybloty, Dec 10 '18 at 18:35

score 0 · Accepted Answer · answered Dec 10 '18 at 18:43

I'm not sure exactly what you plan to do with all the averages, so I'll toss this out there and see if it's helpful. The following code will calculate a bunch of weighted averages, one per unique value of r_vals and store them in a dictionary(which is then printed out).

def calc_weighted_average (r_vals, z_vals, Tgas_vals, n_vals):
    weighted_vals = {} #new variable to store rval=>weighted ave.

    for r in np.unique(r_vals):  
        mask = r_vals == r  # I think yours was backwards
        weighted_vals[r] = np.average(Tgas_vals[mask]*n_vals[mask])

    return weighted_vals

weighted_averages = calc_weighted_average (r_vals, z_vals, Tgas_vals, n_vals)

for rval in weighted_averages:
    print ('%i : %0.4f' % (rval, weighted_averages[rval])) #assuming rval is integer

alternatively, you may want to factor in "z_vals" in somehow. Your question was not clear in this.

Python - masking in a for loop?

1 Answers1