Coverage for ase / visualize / plot.py: 90.24%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING 

4 

5from ase.io.utils import PlottingVariables, make_patch_list 

6 

7if TYPE_CHECKING: 

8 from matplotlib.axes import Axes 

9 

10 from ase.atoms import Atoms 

11 

12 

13class Matplotlib(PlottingVariables): 

14 def __init__( 

15 self, 

16 atoms, 

17 ax, 

18 rotation='', 

19 radii=None, 

20 colors=None, 

21 scale=1, 

22 offset=(0, 0), 

23 **parameters, 

24 ): 

25 super().__init__( 

26 atoms, 

27 rotation=rotation, 

28 radii=radii, 

29 colors=colors, 

30 scale=scale, 

31 extra_offset=offset, 

32 **parameters, 

33 ) 

34 

35 self.ax = ax 

36 self.figure = ax.figure 

37 self.ax.set_aspect('equal') 

38 

39 def write(self): 

40 self.write_body() 

41 self.ax.set_xlim(0, self.w) 

42 self.ax.set_ylim(0, self.h) 

43 

44 def write_body(self): 

45 patch_list = make_patch_list(self) 

46 for patch in patch_list: 

47 self.ax.add_patch(patch) 

48 

49 

50def animate( 

51 images, 

52 ax=None, 

53 interval=200, # in ms; same default value as in FuncAnimation 

54 save_count=None, # ignored as of 2023 with newer matplotlib 

55 **parameters, 

56): 

57 """Convert sequence of atoms objects into Matplotlib animation. 

58 

59 Each image is generated using plot_atoms(). Additional parameters 

60 are passed to this function.""" 

61 import matplotlib.pyplot as plt 

62 from matplotlib.animation import FuncAnimation 

63 

64 if ax is None: 

65 ax = plt.gca() 

66 

67 fig = ax.get_figure() 

68 

69 def drawimage(atoms): 

70 ax.clear() 

71 ax.axis('off') 

72 plot_atoms(atoms, ax=ax, **parameters) 

73 

74 animation = FuncAnimation( 

75 fig, drawimage, frames=images, init_func=lambda: None, interval=interval 

76 ) 

77 return animation 

78 

79 

80def plot_atoms(atoms: Atoms, ax: Axes | None = None, **kwargs) -> Axes: 

81 """Plot an Atoms object in a matplotlib subplot. 

82 

83 Axis decorations will always be removed. 

84 

85 For finer control of plot appearance or composition of custom plot 

86 arrangements, consider using plot_atoms_raw() to plot directly to given 

87 Matplotlib Axes. 

88 

89 Parameters 

90 ---------- 

91 atoms : Atoms object 

92 ax : Matplotlib subplot object 

93 rotation : str, optional 

94 In degrees. In the form '10x,20y,30z' 

95 show_unit_cell : int, optional, default 2 

96 Draw the unit cell as dashed lines depending on value: 

97 0: Don't 

98 1: Do 

99 2: Do, making sure cell is visible 

100 radii : float, optional 

101 The radii of the atoms 

102 colors : list of strings, optional 

103 Color of the atoms, must be the same length as 

104 the number of atoms in the atoms object. 

105 scale : float, optional 

106 Scaling of the plotted atoms and lines. 

107 offset : tuple (float, float), optional 

108 Offset of the plotted atoms and lines. 

109 

110 """ 

111 import matplotlib.pyplot as plt 

112 

113 if ax is None: 

114 _, ax = plt.subplots() 

115 

116 if isinstance(atoms, list): 

117 assert len(atoms) == 1 

118 atoms = atoms[0] 

119 

120 plot_atoms_raw(atoms, ax, **kwargs) 

121 ax.set_axis_off() 

122 

123 return ax 

124 

125 

126def plot_atoms_raw(atoms: Atoms, ax: Axes, **kwargs) -> None: 

127 """Plot Atoms to a matplotlib subplot without additional formatting. 

128 

129 Compared to plot_atoms this is more "barebones". To reproduce the 

130 additional features of plot_atoms: 

131 

132 - create a figure and axes if necessary using plt.subplots() 

133 

134 - disable decorations by calling ax.set_axis_off() 

135 

136 - set matplotlib axis properties directly by calling ax.set() 

137 

138 Parameters 

139 ---------- 

140 atoms : Atoms object 

141 ax : Matplotlib subplot object 

142 rotation : str, optional 

143 In degrees. In the form '10x,20y,30z' 

144 show_unit_cell : int, optional, default 2 

145 Draw the unit cell as dashed lines depending on value: 

146 0: Don't 

147 1: Do 

148 2: Do, making sure cell is visible 

149 radii : float, optional 

150 The radii of the atoms 

151 colors : list of strings, optional 

152 Color of the atoms, must be the same length as 

153 the number of atoms in the atoms object. 

154 scale : float, optional 

155 Scaling of the plotted atoms and lines. 

156 offset : tuple (float, float), optional 

157 Offset of the plotted atoms and lines. 

158 

159 """ 

160 Matplotlib(atoms, ax, **kwargs).write()