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?

Answer 1

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.