nl_grid = np.mgrid[0:lmax+1, 0:nmax+1]

plt.figure(figsize=(5,4))
plt.scatter(nl_grid[0].ravel(), nl_grid[1].ravel(),
            c=np.abs(S[:,:,0].ravel()), norm=mpl.colors.LogNorm(),
            cmap='viridis', s=80)
plt.xlabel('$n$')
plt.ylabel('$l$')
plt.colorbar()
plt.tight_layout()