CODEWITHSUNDEEP
pythonpython-3.xnetworkxphysics
Sid
Sid

Reputation: 65

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

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.

Upvotes: 4

Views: 2655

Answers (2)

Sid
Sid

Reputation: 65

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()

Resultant Plot

Upvotes: 0

Jan Joswig
Jan Joswig

Reputation: 733

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)

enter image description here

Upvotes: 4

Related Questions