I am trying to understand the dynamics of two bosonic particles on N-lattice sites of a Bose-Hubbard Model through the plot prob density of the two particles as a function of sites i, j (where they are located when their initial wavefunction is evolved over time).
I have two input lists X, Y – they store the sites where the two particles are located. Output List prob – contains the prob of finding the two particles at sites (X, Y).
I am expected to get a plot as shown below – X, Y axis represent the sites for N = 20 lattice sites where the two particles are located and the colorbar represents the probability with which these particles can be found at the sites i and j and this plot is particularly done for chemical potential 5.1.
(https://i.sstatic.net/ZiPKNdmS.png)

I went through previous stackoverflow posts related to the same issue I am facing, I tried the solutions mentioned there however, they did not work out for me. Based on what I understood from these posts, for heatmap, I had to convert all these lists to an array, and find out the unique elements, dim for x and y and create datapoints for the map. And, then assign prob values (np.abs(prob))**2 to the points heatmap[x, y]. I want to know if this is the correct way of doing it because the plots are not matching from that of the paper (fig 2 of the paper – https://www.science.org/doi/10.1126/science.1260364) (and, the code for constructing the hamiltonian matrix seems to be correct since it matches with the pen-and-paper calculations)

Here’s the code:

X = np.array(X)
Y = np.array(Y)
probabilities = np.array(probabilities)

plt.figure(figsize=[7, 5])
p = plt.get_cmap('Blues')

x_unique = np.unique(X)
y_unique = np.unique(Y)
heatmap_data = np.zeros((len(y_unique), len(x_unique)))

x_index = {value: idx for idx, value in enumerate(x_unique)}
y_index = {value: idx for idx, value in enumerate(y_unique)}

for (x, y, prob) in zip(X, Y, probabilities):
    i = y_index[y]
    j = x_index[x]
    heatmap_data[i, j] = (np.abs(prob))**2

plt.figure(figsize=(10, 8))
sns.heatmap(heatmap_data, cmap=p)
plt.xlabel('Site i')
plt.ylabel('Site j')
mu = params['mu']
plt.title(f'Probability as a function of Sites i and j, Δ={mu}')
plt.show()