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.
Related to: Trying to convert adjacency list to adjacency matrix in Python
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)
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
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.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.