Coverage for /builds/kinetik161/ase/ase/visualize/plot.py: 88.89%

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=None, # ignored as of 2023 with newer matplotlib 

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 def drawimage(atoms): 

45 ax.clear() 

46 ax.axis('off') 

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

48 

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

50 init_func=lambda: None, 

51 interval=interval) 

52 return animation 

53 

54 

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

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

57 

58 Parameters 

59 ---------- 

60 atoms : Atoms object 

61 ax : Matplotlib subplot object 

62 rotation : str, optional 

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

64 show_unit_cell : int, optional, default 2 

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

66 0: Don't 

67 1: Do 

68 2: Do, making sure cell is visible 

69 radii : float, optional 

70 The radii of the atoms 

71 colors : list of strings, optional 

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

73 the number of atoms in the atoms object. 

74 scale : float, optional 

75 Scaling of the plotted atoms and lines. 

76 offset : tuple (float, float), optional 

77 Offset of the plotted atoms and lines. 

78 """ 

79 if isinstance(atoms, list): 

80 assert len(atoms) == 1 

81 atoms = atoms[0] 

82 

83 import matplotlib.pyplot as plt 

84 if ax is None: 

85 ax = plt.gca() 

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

87 return ax