Add atomic images to a pyplot figure

by Andrew Peterson

It can be nice to add atomic images directly to a figure made in pyplot (i.e., matplotlib). This is pretty simple to do using the tools of ASE. Here we just show a very simple example of plotting a row of atoms with an adsorbate on top.

import numpy as np
from matplotlib import pyplot
from ase.lattice.surface import fcc111
from ase import Atom
from ase.data import covalent_radii as radii
from ase.data.colors import jmol_colors
from matplotlib.patches import Circle

fig, ax = pyplot.subplots()

# Make the atoms; rotate so they are viewable with no z coordinates.
atoms = fcc111('Pt', (3, 1, 1), vacuum=20, a=4.0)
atoms.append(Atom('N', atoms[1].position + (0., 0., 2.1)))
atoms.rotate('x', -np.pi/2., rotate_cell=True)

# Add the atoms to the plot as circles.
for atom in atoms:
    color = jmol_colors[atom.number]
    radius = radii[atom.number]
    circle = Circle((atom.x, atom.y), radius, facecolor=color,
                            edgecolor='k', linewidth=1)
    ax.add_patch(circle)

# Enforce the circles are round (equal x and y scales) and turn off
# tickmarks.
ax.axis('equal')
ax.set_xticks([])
ax.set_yticks([])
ax.axis('off')

fig.savefig('out.png')

Most of the script above should be pretty self-explanatory, provided you are already comfortable with matplotlib and ASE. The atoms are added to the plot as Circle patches. The trick is to pull the atomic radii and colors from ASE, which are done with the covalent_radii and jmol_colors imports, respectively. After this, the atoms show up just like they would in ASE gui, except on an axes canvas, so you can add any standard features from matplotlib on top of this figure. The output of this simple script is shown below.