功能变化不大。 面具。 蟒蛇

Question

I have the following function:

def delta(r, dr):
   res = np.zeros(r.shape)
   mask1 = (r >= 0.5*dr) & (r <= 1.5*dr)
   res[mask1] = (5-3*np.abs(r[mask1])/dr \
    - np.sqrt(-3*(1-np.abs(r[mask1])/dr)**2+1))/(6*dr)
   mask2 = np.logical_not(mask1) & (r <= 0.5*dr)
   res[mask2] = (1+np.sqrt(-3*(r[mask2]/dr)**2+1))/(3*dr)
   return res

Where r is a numpy.array of size (shape[0],shape[1]) and dr is a single value. I want to modify the function to make dr also an array of the same size as r and for each value of r take the analogous value from dr .

For example r[0,0] goes with dr[0,0] , r[0,1] with dr[0,1] and so on. Any ideas?

Answer 1

You could multiply the 2D mask with the input array, which in effect is masking and thus perform the computations resulting in a 2D array instead of 1D array with boolean indexing as done so far. The only difference would be setting values into the output array, for which you need to mask both the array to be set and the 2D computed array from which values would be selected.

The implementation would look like this -

# Initialize output array   
res = np.zeros(r.shape)

# Get mask1 and compute values for all elements and use the mask to set only
# TRUE positions with the computed values
mask1 = (r >= 0.5*dr) & (r <= 1.5*dr)
V1 = (5-3*np.abs(r*mask1)/dr - np.sqrt(-3*(1-np.abs(r*mask1)/dr)**2+1))/(6*dr)
res[mask1] = V1[mask1]

# Similarly for mask2 and the computations with that mask
mask2 = np.logical_not(mask1) & (r <= 0.5*dr)
V2 = (1+np.sqrt(-3*(r*mask2/dr)**2+1))/(3*dr)
res[mask2] = V2[mask2]