Best way to plot a 2D square lattice of atoms in Python?

Question

Summary of the problem: I'm working on a Physics problem and I want to plot a 2D lattice of atoms, where the nodes have been connected using arrows as shown here 2D lattice figure.

What I've tried: I've tried using grid_2d_graph from NetworkX, taking help from this answer but wasn't able to get it to work as I wanted it to. The code I used was as follows:

G = nx.grid_2d_graph(4,4)
pos = dict( (n, n) for n in G.nodes() )
nx.draw_networkx(G, pos=pos)
plt.axis('off')
plt.show()

This produced the following image, which wasn't exactly what I had in mind.

And what determines the direction of the arrows? – yatu Jun 02 '20 at 07:28 — yatu, Jun 02 '20 at 07:28

score 4 · Accepted Answer · answered Jun 03 '20 at 05:12

Here is a simplistic approach using matplotlib's arrow. It requires, that the grid (graph) is represented as a dictionary where the keys are point (node) coordinates on the grid and the values are neighbouring points to which an outgoing arrow (edge) should be drawn. When the size of your grid changes, you will probably need to play around with w, h and such to control the size of the plotted elements.

grid = {  # point(x, y), outgoing connections [points] 
    (0, 0): [(0, 1), (1, 0)],
    (0, 1): [],
    (1, 0): [],
    (1, 1): [(1, 0), (0, 1)]
}

w = 0.005  # Errorwidth
h = 0.05   # Errorhead width

fig, ax = plt.subplots()
for point, connections in grid.items():
    for outgoing in connections:
        dx = outgoing[0] - point[0]
        dy = outgoing[1] - point[1]
        ax.arrow(point[0], point[1],
                 dx / 2, dy / 2,
                 width=w, head_width=h,
                 facecolor="k",
                 zorder=0)
        ax.arrow(point[0] + dx / 2,
                 point[1] + dy / 2,
                 dx / 2, dy / 2,
                 width=w, head_width=0,
                 facecolor="k",
                 zorder=0)
    ax.plot(*point,
            marker="o", markersize=10, markeredgecolor="k",
            markerfacecolor="red",
            zorder=1)

score 0 · Answer 2 · answered Jun 03 '20 at 09:07

I was trying to use quiver plots when I came up on Jan's answer. I've modified his code to work with quiver plots

def plot_vector(p1,p2):
    p1 = np.array(p1)
    p2 = np.array(p2)
    dp = p2-p1
    plt.quiver(p1[0], p1[1], dp[0], dp[1],angles='xy', scale_units='xy', scale=1, headwidth = 5, headlength = 7)


grid = {  # point(x, y), outgoing connections [points] 
    (0, 0): [(0, 1), (1, 0)],
    (0, 1): [],
    (1, 0): [],
    (1, 1): [(1, 0), (0, 1)]
}


fig, ax = plt.subplots()
for point, connections in grid.items():
    for outgoing in connections:
        plot_vector(point,outgoing)
    plt.plot(*point,
            marker="o", markersize=10, markeredgecolor="k",
            markerfacecolor="red",
            zorder=1)
plt.show()

Best way to plot a 2D square lattice of atoms in Python?

2 Answers2