import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D   # registers the 3D projection
import matplotlib.cm as cm
import matplotlib.colors as colors

w_vals = np.arange(1, 25)
d_vals = np.arange(1, 25)

W, D = np.meshgrid(w_vals, d_vals)   # shape (len(d_vals), len(w_vals))
heights = (W - 1) * (D - 1) + 1      # same shape

x = W.ravel()
y = D.ravel()
z = np.zeros_like(x)   # base of bars
dx = dy = 0.6          # bar width in x and y
dz = heights.ravel()

norm = colors.Normalize(vmin=dz.min(), vmax=dz.max())
colors_map = cm.viridis(norm(dz))

fig = plt.figure(figsize=(10, 6))
ax = fig.add_subplot(111, projection='3d')

ax.bar3d(x - dx/2, y - dy/2, z, dx, dy, dz, shade=True, color=colors_map, zsort='average')

ax.set_xlabel('Width')
ax.set_ylabel('Depth')
ax.set_zlabel('#Atomic Models')
ax.set_xticks(w_vals)
ax.set_yticks(d_vals)
ax.set_title("#Atomic Models in function of width and depth")
ax.view_init(elev=30, azim=-135)   # example values

# colorbar
mappable = cm.ScalarMappable(norm=norm, cmap=cm.viridis)
mappable.set_array(dz)
cbar = fig.colorbar(mappable, ax=ax, shrink=0.6, pad=0.1)
cbar.set_label('value')

plt.tight_layout()
plt.show()
# Optionally save:
# fig.savefig("3d_bars_matplotlib.png", dpi=200)