Fancy indexing calculation of adjacency matrix from adjacency list-CodePudding

Problem:

I want to calculate at several times the adjacency matrix A_ij given the adjacency list E_ij, where E_ij[t,i] = j gives the edge from i to j at time t.

I can do it with the following code:

import numpy as np

nTimes = 100
nParticles = 10
A_ij = np.full((nTimes, nParticles, nParticles), False)
E_ij = np.random.randint(0, 9, (100, 10))

for t in range(nTimes):
    for i in range(nParticles):
        A_ij[t, i, E_ij[t,i]] = True

Question:

How can I do it in a vectorized way, either with fancy indexing or using numpy functions such as np.take_along_axis?

What I tried:

I expected this to work:

A_ij[:,np.arange(nParticles)[None,:,None], E_ij[:,None,np.arange(nParticles)]] = True

But it does not.

CodePudding user response：

I think this might work:

import numpy as np

nTimes = 100
nParticles = 10
A_ij = np.full((nTimes, nParticles, nParticles), False)
E_ij = np.random.randint(0, 9, (100, 10))

np.put_along_axis(A_ij, E_ij[..., None], True, axis=2)

CodePudding user response：

Another way to do it, and close to what you actually tried at the end would be something like:

i, j = np.mgrid[:nTimes, :nParticles]
A_ij[i, j, E_ij] = True

But the accepted answer is definitely the better way to go about the problem, no need to construct indices.

CodePudding user response：

In case it may help other people, I also found a way to do fancy indexing in this problem, but @Chrysophylaxs answer was faster and simpler (I guess I was confused with the indices and I could not think about it). I also add @Mercury answer for comparison.

Code:

import numpy as np
import matplotlib.pyplot as plt
import time


nTimes = 1000000
nParticles = 10
A_ij1 = np.full((nTimes, nParticles, nParticles), False)
A_ij2 = np.full((nTimes, nParticles, nParticles), False)
A_ij3 = np.full((nTimes, nParticles, nParticles), False)
A_ij4 = np.full((nTimes, nParticles, nParticles), False)


E_ij = np.random.randint(0, 9, (nTimes, 10))

start_time = time.time()
for t in range(nTimes):
    for i in range(nParticles):
        A_ij1[t, i, E_ij[t,i]] = True
print("Loop: %s s" % (time.time() - start_time))

        
start_time = time.time()
A_ij2[np.arange(nTimes)[:,None],np.arange(nParticles)[None,:], E_ij[np.arange(nTimes)[:,None],np.arange(nParticles)[None,:]]] = True
print("Fancy indexing: %s s" % (time.time() - start_time))

start_time = time.time()
np.put_along_axis(A_ij3, E_ij[..., None], True, axis=2)
print("Put along axis: %s s" % (time.time() - start_time))

start_time = time.time()
i, j = np.mgrid[:nTimes, :nParticles]
A_ij4[i, j, E_ij] = True
print("mgrid: %s s" % (time.time() - start_time))


print(np.allclose(A_ij1, A_ij2))
print(np.allclose(A_ij1, A_ij3))
print(np.allclose(A_ij1, A_ij4))

Output:

Loop: 2.5006823539733887 s
Fancy indexing: 0.11996173858642578 s
Put along axis: 0.0814671516418457 s
mgrid: 0.19223332405090332 s
True
True
True