我想你必须使用某种形式的循环来实现这一点,但也许这会有所帮助:
import numpy as np
import matplotlib.pyplot as plt
L = 50
N = 30
W = 20
n_steps = 1000
# Initialize all walkers on the left side
wn0 = np.arange(W)[:, np.newaxis].repeat(N, axis=1)
# Set up the plot
fig, ax = plt.subplots()
worlds = np.zeros((N, L))
worlds[np.arange(N)[np.newaxis, :], wn0] = np.arange(W)[:, np.newaxis]
h = ax.imshow(worlds, cmap='gray_r') # cmap='tab20')
ax.set_xlabel('Distance in 1D World')
ax.set_ylabel('Ensemble of Worlds')
for _ in range(n_steps):
r = np.where(np.random.random(wn0.shape) < 0.5, 1, -1)
wn1 = wn0 + r
wn1 = np.clip(wn1, 0, L-1)
# Case 1
rest_mat = np.zeros_like(wn0, dtype=bool)
for i in range(W):
for j in range(i+1, W):
rest_mat[[[i], [j]], np.logical_and(wn0[i] == wn1[j], wn1[i] == wn0[j])] = True
wn1[rest_mat] = wn0[rest_mat]
# Case 2, go from 0->W and than from W->0 to make sure all duplicates are gone
for i in np.hstack((range(W), range(W)[-2::-1])):
temp = (wn1[i] == wn1).sum(axis=0) > 1
wn1[i, temp] = wn0[i, temp]
# for wn1_j in wn1.T: Check if there are collisions
# assert len(np.unique(wn1_j)) == W
wn0 = wn1
worlds = np.zeros((N, L))
worlds[np.arange(N)[np.newaxis, :], wn1] = np.arange(W)[:, np.newaxis]
h.set_data(worlds)
plt.pause(0.1)
