CODEWITHSUNDEEP
pythonarraysnumpyindexing
JP K.
JP K.

Reputation: 1301

python - numpy fancy broadcasting for special case riddle

I want to do some forces calculations between vertices and because the forces are symmetrical I have a list of vertice-pairs that need those forces added. I am sure it's possible with fancy indexing, but I really just can get it to work with a slow python for-loop. for symmetric reasons, the right-hand side of the index array needs a negative sign when adding the forces.

consider you have the vertice index array:

>>> I = np.array([[0,1],[1,2],[2,0]])
I = [[0 1]
     [1 2]
     [2 0]]

and the x,y forces array for each pair:

>>> F = np.array([[3,6],[4,7],[5,8]])
F = [[3 6]
     [4 7]
     [5 8]]

the wanted operation could be described as:

"vertice #0 sums the force vectors (3,6) and (-5,-8),
 vertice #1 sums the force vectors (-3,-6) and (4,7),
 vertice #2 sums the force vectors (-4,-7) and (5,8)"

Desired results:

     [ 3  6 ]   [ 0  0 ]   [-5 -8 ]   [-2 -2 ] //resulting force Vertice #0
A =  [-3 -6 ] + [ 4  7 ] + [ 0  0 ] = [ 1  1 ] //resulting force Vertice #1
     [ 0  0 ]   [-4 -7 ]   [ 5  8 ]   [ 1  1 ] //resulting force Vertice #2

edit:

my ugly for-loop solution:

import numpy as np

I = np.array([[0,1],[1,2],[2,0]])
F = np.array([[3,6],[4,7],[5,8]])
A = np.zeros((3,2))

A_x = np.zeros((3,2))
A_y = np.zeros((3,2))

for row in range(0,len(F)):


    A_x[I[row][0],0]= F[row][0]
    A_x[I[row][1],1]= -F[row][0] 

    A_y[I[row][0],0]= F[row][1]
    A_y[I[row][1],1]= -F[row][1] 


A = np.hstack((np.sum(A_x,axis=1).reshape((3,1)),np.sum(A_y,axis=1).reshape((3,1)))) 

print(A)

A=  [[-2. -2.]
     [ 1.  1.]
     [ 1.  1.]]

Upvotes: 1

Views: 117

Answers (3)

Paul Panzer
Paul Panzer

Reputation: 53029

You can preallocate an array to hold the shuffled forces and then use the index like so:

>>> N = I.max() + 1
>>> out = np.zeros((N, 2, 2), F.dtype)
>>> out[I, [1, 0]] = F[:, None, :]
>>> np.diff(out, axis=1).squeeze()
array([[-2, -2],
       [ 1,  1],
       [ 1,  1]])

or, equivalently,

>>> out = np.zeros((2, N, 2), F.dtype)
>>> out[[[1], [0]], I.T] = F
>>> np.diff(out, axis=0).squeeze()
array([[-2, -2],
       [ 1,  1],
       [ 1,  1]])

Upvotes: 0

Nils Werner
Nils Werner

Reputation: 36765

Your current "push-style" interpretation of I is

For row-index k in I, take the forces from F[k] and add/subtract them to out[I[k], :]

I = np.array([[0,1],[1,2],[2,0]])
out = numpy.zeros_like(F)
for k, d in enumerate(I):
    out[d[0], :] += F[k]
    out[d[1], :] -= F[k]
out
# array([[-2, -2],
#        [ 1,  1],
#        [ 1,  1]])

However you can also change the meaning of I on its head and make it "pull-style", so it says

For row-index k in I, set vertex out[k] to be the difference of F[I[k]]

I = np.array([[0,2],[1,0],[2,1]])
out = numpy.zeros_like(F)
for k, d in enumerate(I):
    out[k, :] = F[d[0], :] - F[d[1], :]
out
# array([[-2, -2],
#        [ 1,  1],
#        [ 1,  1]])

In which case the operation simplifies quite easily to mere fancy indexing:

out = F[I[:, 0], :] - F[I[:, 1], :]
# array([[-2, -2],
#        [ 1,  1],
#        [ 1,  1]])

Upvotes: 1

WolfgangK
WolfgangK

Reputation: 993

The way I understand the question, the values in the I array represent the vortex number, or the name of the vortex. They are not an actual positional index. Based on this thought, I have a different solution that uses the original I array. It does not quite come without loops, but should be OK for a reasonable number of vertices:

I = np.array([[0,1],[1,2],[2,0]])   
F = np.array([[3,6],[4,7],[5,8]])

pos = I[:, 0]
neg = I[:, 1]
A = np.zeros_like(F)

unique = np.unique(I)
for i, vortex_number in enumerate(unique):
    A[i] = F[np.where(pos==vortex_number)] - F[np.where(neg==vortex_number)]

# produces the expected result
# [[-2 -2]
#  [ 1  1]
#  [ 1  1]]

Maybe this loop can also be replaced by some numpy magic.

Upvotes: 0

Related Questions