1from ase.io.utils import PlottingVariables, make_patch_list 

2 

3 

4class Matplotlib(PlottingVariables): 

5 def __init__(self, atoms, ax, 

6 rotation='', radii=None, 

7 colors=None, scale=1, offset=(0, 0), **parameters): 

8 PlottingVariables.__init__( 

9 self, atoms, rotation=rotation, 

10 radii=radii, colors=colors, scale=scale, 

11 extra_offset=offset, **parameters) 

12 

13 self.ax = ax 

14 self.figure = ax.figure 

15 self.ax.set_aspect('equal') 

16 

17 def write(self): 

18 self.write_body() 

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

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

21 

22 def write_body(self): 

23 patch_list = make_patch_list(self) 

24 for patch in patch_list: 

25 self.ax.add_patch(patch) 

26 

27 

28def animate(images, ax=None, 

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

30 save_count=100, 

31 **parameters): 

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

33 

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

35 are passed to this function.""" 

36 import matplotlib.pyplot as plt 

37 from matplotlib.animation import FuncAnimation 

38 

39 if ax is None: 

40 ax = plt.gca() 

41 

42 fig = ax.get_figure() 

43 

44 nframes = [0] 

45 

46 def drawimage(atoms): 

47 ax.clear() 

48 ax.axis('off') 

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

50 nframes[0] += 1 

51 # Animation will stop without warning if we don't have len(). 

52 # Write a warning if we may be missing frames: 

53 if not hasattr(images, '__len__') and nframes[0] == save_count: 

54 import warnings 

55 warnings.warn('Number of frames reached animation savecount {}; ' 

56 'some frames may not be saved.' 

57 .format(save_count)) 

58 

59 animation = FuncAnimation(fig, drawimage, frames=images, 

60 init_func=lambda: None, 

61 save_count=save_count, 

62 interval=interval) 

63 return animation 

64 

65 

66def plot_atoms(atoms, ax=None, **parameters): 

67 """Plot an atoms object in a matplotlib subplot. 

68 

69 Parameters 

70 ---------- 

71 atoms : Atoms object 

72 ax : Matplotlib subplot object 

73 rotation : str, optional 

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

75 show_unit_cell : int, optional, default 2 

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

77 0: Don't 

78 1: Do 

79 2: Do, making sure cell is visible 

80 radii : float, optional 

81 The radii of the atoms 

82 colors : list of strings, optional 

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

84 the number of atoms in the atoms object. 

85 scale : float, optional 

86 Scaling of the plotted atoms and lines. 

87 offset : tuple (float, float), optional 

88 Offset of the plotted atoms and lines. 

89 """ 

90 if isinstance(atoms, list): 

91 assert len(atoms) == 1 

92 atoms = atoms[0] 

93 

94 import matplotlib.pyplot as plt 

95 if ax is None: 

96 ax = plt.gca() 

97 Matplotlib(atoms, ax, **parameters).write() 

98 return ax