Skip to content

atomlib.visualize

Visualization of atomic structures. Useful for debugging.

BackendName module-attribute 

BackendName: TypeAlias = Literal['mpl', 'ase']

AtomStyle module-attribute 

AtomStyle: TypeAlias = Literal[
    "spacefill", "ballstick", "small"
]

AtomImage

Bases: ABC

Source code in atomlib/visualize/__init__.py 
33
34
35
36
class AtomImage(ABC):
    @abstractmethod
    def save(self, f: FileOrPath):
        ...

save abstractmethod 

save(f: FileOrPath)
Source code in atomlib/visualize/__init__.py 
34
35
36
@abstractmethod
def save(self, f: FileOrPath):
    ...

AtomImageMpl

Bases: Figure, AtomImage

Source code in atomlib/visualize/__init__.py 
76
77
78
79
80
81
82
83
84
85
class AtomImageMpl(Figure, AtomImage):
    def __new__(cls, fig: Figure):
        fig.__class__ = cls
        return fig

    def __init__(self, fig: Figure):
        ...

    def save(self, f: FileOrPath):
        return self.savefig(f)  # type: ignore

save 

save(f: FileOrPath)
Source code in atomlib/visualize/__init__.py 
84
85
def save(self, f: FileOrPath):
    return self.savefig(f)  # type: ignore

show_atoms_3d 

show_atoms_3d(
    atoms: HasAtoms,
    *,
    zone: Optional[VecLike] = None,
    plane: Optional[VecLike] = None,
    backend: BackendName = "mpl",
    style: AtomStyle = "small",
    **kwargs: Any
) -> AtomImage

Show atoms on a 3D plot, using backend backend (defaults to matplotlib).

Source code in atomlib/visualize/__init__.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def show_atoms_3d(atoms: HasAtoms, *,
                  zone: t.Optional[VecLike] = None,
                  plane: t.Optional[VecLike] = None,
                  backend: BackendName = 'mpl',
                  style: AtomStyle = 'small', **kwargs: t.Any) -> AtomImage:
    """
    Show `atoms` on a 3D plot, using backend `backend` (defaults to matplotlib).
    """

    backend = t.cast(BackendName, backend.lower())
    if backend == 'mpl':
        return show_atoms_mpl_3d(atoms, zone=zone, plane=plane, style=style, **kwargs)
    elif backend == 'ase':
        raise NotImplementedError()

    raise ValueError(f"Unknown backend '{backend}'")

show_atoms_2d 

show_atoms_2d(
    atoms: HasAtoms,
    *,
    zone: Optional[VecLike] = None,
    plane: Optional[VecLike] = None,
    horz: Optional[VecLike] = None,
    backend: BackendName = "mpl",
    style: AtomStyle = "small",
    **kwargs: Any
) -> AtomImage

Show atoms on a 2D plot, using backend backend (defaults to matplotlib).

Source code in atomlib/visualize/__init__.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def show_atoms_2d(atoms: HasAtoms, *,
                  zone: t.Optional[VecLike] = None,
                  plane: t.Optional[VecLike] = None,
                  horz: t.Optional[VecLike] = None,
                  backend: BackendName = 'mpl',
                  style: AtomStyle = 'small', **kwargs: t.Any) -> AtomImage:
    """
    Show `atoms` on a 2D plot, using backend `backend` (defaults to matplotlib).
    """

    backend = t.cast(BackendName, backend.lower())
    if backend == 'mpl':
        return show_atoms_mpl_2d(atoms, zone=zone, plane=plane, horz=horz, style=style, **kwargs)
    elif backend == 'ase':
        raise NotImplementedError()

    raise ValueError(f"Unknown backend '{backend}'")

get_elem_color 

get_elem_color(elem: int) -> List[int]

Get the color to use for element elem.

Source code in atomlib/visualize/__init__.py 
 99
100
101
102
def get_elem_color(elem: int) -> t.List[int]:
    """Get the color to use for element `elem`."""
    # grey fallback
    return _ELEM_MAP.get(elem, [80, 80, 80])

get_zone 

get_zone(
    atoms: HasAtoms,
    zone: Optional[VecLike] = None,
    plane: Optional[VecLike] = None,
    default: Optional[VecLike] = None,
) -> NDArray[float64]

Get the zone axis corresponding to the arguments zone, plane, and default.

Source code in atomlib/visualize/__init__.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def get_zone(atoms: HasAtoms, zone: t.Optional[VecLike] = None,
             plane: t.Optional[VecLike] = None,
             default: t.Optional[VecLike] = None) -> NDArray[numpy.float64]:
    """
    Get the zone axis corresponding to the arguments `zone`, `plane`, and `default`.
    """

    if zone is not None and plane is not None:
        raise ValueError("'zone' and 'plane' can't both be specified.")
    if plane is not None:
        if isinstance(atoms, AtomCell) and not atoms.is_orthogonal():
            # convert plane into zone
            raise NotImplementedError()
        zone = plane
    if zone is not None:
        return numpy.broadcast_to(zone, 3).astype(numpy.float64)
    if default is not None:
        return numpy.broadcast_to(default, 3).astype(numpy.float64)
    return numpy.array([0., 0., 1.], dtype=numpy.float64)

get_plot_radii 

get_plot_radii(
    atoms: HasAtoms,
    min_r: Optional[float] = 1.0,
    style: AtomStyle = "small",
) -> NDArray[float64]

Get the radii to use for each atom in atoms.

Source code in atomlib/visualize/__init__.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def get_plot_radii(atoms: HasAtoms, min_r: t.Optional[float] = 1.0, style: AtomStyle = 'small') -> NDArray[numpy.float64]:
    """Get the radii to use for each atom in `atoms`."""
    radii = get_radius(atoms['elem']).to_numpy()
    if style == 'small':
        radii = radii * 0.6
    elif style == 'ballstick':
        radii = radii * 0.5
    elif style == 'spacefill':
        radii = radii * 1.0
    else:
        raise ValueError(f"Unknown atom style '{style}'. Expected 'spacefill', 'ballstick', or 'small'.")
    if min_r is not None:
        return numpy.maximum(min_r, radii)
    return radii

get_azim_elev 

get_azim_elev(zone: VecLike) -> Tuple[float, float]

Get the azimuth and elevation corresponding to the zone zone.

Source code in atomlib/visualize/__init__.py 
142
143
144
145
146
def get_azim_elev(zone: VecLike) -> t.Tuple[float, float]:
    """Get the azimuth and elevation corresponding to the zone `zone`."""
    (a, b, c) = -to_vec3(zone)  # look down zone
    # todo: aren't these just arctan2s?
    return (numpy.angle(a + b*1.j, deg=True), numpy.angle(numpy.sqrt(a**2 + b**2) + c*1.j, deg=True))  # type: ignore

show_atoms_mpl_3d 

show_atoms_mpl_3d(
    atoms: HasAtoms,
    *,
    fig: Optional[Figure] = None,
    zone: Optional[VecLike] = None,
    plane: Optional[VecLike] = None,
    min_r: Optional[float] = 1.0,
    style: AtomStyle = "small"
) -> AtomImageMpl

Show atoms on a 3D plot using matplotlib.

Source code in atomlib/visualize/__init__.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
def show_atoms_mpl_3d(atoms: HasAtoms, *, fig: t.Optional[Figure] = None,
                      zone: t.Optional[VecLike] = None,
                      plane: t.Optional[VecLike] = None,
                      min_r: t.Optional[float] = 1.0,
                      style: AtomStyle = 'small') -> AtomImageMpl:
    """Show `atoms` on a 3D plot using matplotlib."""

    fig = AtomImageMpl(fig or pyplot.figure())  # type: ignore

    zone = get_zone(atoms, zone, plane, [1., 2., 4.])
    (azim, elev) = get_azim_elev(zone)

    rect = (0., 0., 1., 1.)
    ax: Axes3D = fig.add_axes(rect, axes_class=Axes3D, proj_type='ortho', azim=azim, elev=elev)  # type: ignore
    ax.grid(False)

    bbox = atoms.bbox().pad(0.2)
    ax.set_xlim3d(bbox.x)  # type: ignore
    ax.set_ylim3d(bbox.y)  # type: ignore
    ax.set_zlim3d(bbox.z)  # type: ignore
    ax.set_box_aspect(bbox.size)

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    frame = atoms.get_atoms('local')
    #radii = get_plot_radii(atoms, min_r=min_r, style=style)
    coords = frame.coords()
    elem_colors = numpy.array(list(map(get_elem_color, frame['elem']))) / 255.
    s = 100

    if isinstance(atoms, HasCell):
        # plot cell corners
        corners = atoms.corners('global')
        faces = [
            numpy.array([
                corners[(val*2**axis + v1*2**((axis+1) % 3) + v2*2**((axis+2) % 3))]
                for (v1, v2) in [(0, 0), (0, 1), (1, 1), (1, 0), (0, 0)]
            ])
            for axis in (0, 1, 2)
            for val in (0, 1)
        ]
        for face in faces:
            ax.plot3D(*split_arr(face, axis=-1), '.-k', alpha=1, markersize=8)

    ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], c=elem_colors, alpha=1, s=s)  # type: ignore

    return t.cast(AtomImageMpl, fig)

show_atoms_mpl_2d 

show_atoms_mpl_2d(
    atoms: HasAtoms,
    *,
    fig: Optional[Figure] = None,
    zone: Optional[VecLike] = None,
    plane: Optional[VecLike] = None,
    horz: Optional[VecLike] = None,
    min_r: Optional[float] = 1.0,
    style: AtomStyle = "small"
) -> AtomImageMpl

Show atoms on a 2D plot using matplotlib.

Source code in atomlib/visualize/__init__.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def show_atoms_mpl_2d(atoms: HasAtoms, *, fig: t.Optional[Figure] = None,
                      zone: t.Optional[VecLike] = None,
                      plane: t.Optional[VecLike] = None,
                      horz: t.Optional[VecLike] = None,
                      min_r: t.Optional[float] = 1.0,
                      style: AtomStyle = 'small') -> AtomImageMpl:
    """Show `atoms` on a 2D plot using matplotlib."""

    zone = get_zone(atoms, zone, plane, [0., 0., 1.])
    fig = AtomImageMpl(fig or pyplot.figure())  # type: ignore

    rect = (0.05, 0.05, 0.95, 0.95)
    ax: Axes = fig.add_axes(rect)
    ax.set_aspect('equal')

    frame = atoms.get_atoms('local')
    coords = frame.coords()
    elem_colors = numpy.array(list(map(get_elem_color, frame['elem']))) / 255.
    radii = get_plot_radii(frame, min_r=min_r, style=style)

    # look down zone
    if horz is None:
        transform = LinearTransform3D.align_to(zone, [0, 0, -1])
    else:
        transform = LinearTransform3D.align_to(zone, [0, 0, -1], horz, [1, 0, 0])
    bbox_2d = transform @ atoms.bbox()
    coords = transform @ coords
    # sort by z-order
    sort = numpy.argsort(coords[..., 2])
    coords = coords[sort]
    elem_colors = elem_colors[sort]
    radii = radii[sort]

    ax.set_xbound(*bbox_2d.x)
    ax.set_ybound(*bbox_2d.y)

    # old plotting method
    # ax.scatter(coords[:, 0], coords[:, 1], c=elem_colors, alpha=1., s=s)

    ax.add_collection(EllipseCollection(
        radii, radii, numpy.zeros_like(radii), units='xy', facecolors=elem_colors, ec='black',
        offsets=coords[..., :2], offset_transform=ax.transData,
    ))  # type: ignore (bad api typing)

    return t.cast(AtomImageMpl, fig)